Problem Statement in English

You’re given a m x n binary matrix matrix. Each cell either has a 0 or 1. Return the number of square submatrices that have all ones.

Approach

We can use dynamic programming to solve this problem. We can maintain a dp matrix where dp[i][j] represents the size of the largest square submatrix with all ones that can be formed using the cell (i, j) as the bottom-right corner.

To fill the dp matrix, we can iterate through each cell in the input matrix. If the cell contains a 1, we can calculate the size of the largest square submatrix that can be formed using that cell as the bottom-right corner by taking the minimum of the three neighboring cells (top, left, and top-left diagonal) in the dp matrix and adding 1 to it.

The number of square submatrices with all ones equals the the max side length of the square submatrix that can be formed using that cell as the bottom-right corner. We can keep a running total of the count of square submatrices as we fill the dp matrix.

And we’re done!

Solution in Python


class Solution:
    def countSquares(self, matrix: List[List[int]]) -> int:
        m = len(matrix)
        n = len(matrix[0])

        res = 0

        dp = [[0] * n for _ in range(m)]

        for i in range(m):
            for j in range(n):

                if not matrix[i][j]:
                    continue

                top = left = diagLeft = 0

                # top
                if 0<=i-1<m:
                    top = dp[i-1][j]
                
                # left
                if 0<=j-1<n:
                    left = dp[i][j-1]

                # diagonally left
                if 0<=i-1<m and 0<=j-1<n:
                    diagLeft = dp[i-1][j-1]

                nei = min(top, left, diagLeft)

                if not nei:
                    dp[i][j] = 1
                    res+=1
                else:
                    dp[i][j] = nei+1
                    res+=dp[i][j]

        return res

Complexity

  • Time: $O(M \times N)$
    Where M and N are the number of rows and columns in the input matrix respectively. We iterate through each cell in the input matrix once.

  • Space: $O(M \times N)$
    Where M and N are the number of rows and columns in the input matrix respectively. We use a dp matrix of the same size as the input matrix to store the size of the largest square submatrix with all ones that can be formed using each cell as the bottom-right corner.

And we are done.