Problem statement

https://binarysearch.com/problems/Count-Square-Submatrices/

Solution

Equal to Leetcode 1277. Count Square Submatrices with All Ones.

Complexity

It is O(nm) for time and space.

Code

class Solution:
    def solve(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