Leetcode 1277. Count Square Submatrices with All Ones

Problem statement

https://leetcode.com/problems/count-square-submatrices-with-all-ones/

Solution

Use dp, where dp[i][j] is number of squares with left upper corner in (i, j).

Complexity

It is O(nm) for time and space.

Code

class Solution:
    def countSquares(self, M):
        m, n, ans = len(M[0]), len(M), 0

        dp = [[0]*(m + 1) for _ in range(n + 1)]
        for i in range(n):
            for j in range(m):
                if M[i][j] == 1:
                    dp[i][j] = min(dp[i-1][j], dp[i][j-1], dp[i-1][j-1]) + 1
                ans += dp[i][j]         
        return ans