Count Square Submatrices with All Ones LeetCode solution in Python

1
# Runtime: 1687 ms (Top 8.20%) | Memory: 16.3 MB (Top 73.07%)
2
class Solution:
3
    def countSquares(self, matrix: List[List[int]]) -> int:
4

5
        m = len(matrix)
6
        n = len(matrix[0])
7

8
        dp = [[0 for _ in range(n)] for _ in range(m)]
9
        total = 0
10

11
        for i in range(m):
12
            for j in range(n):
13

14
                if i == 0:
15
                    dp[i][j] = matrix[0][j]
16

17
                elif j == 0:
18
                    dp[i][j] = matrix[i][0]
19

20
                else:
21
                    if matrix[i][j] == 1:
22
                        dp[i][j] = 1 + min(dp[i][j - 1], dp[i - 1][j - 1], dp[i - 1][j])
23

24
                total += dp[i][j]
25

26
        return total