Leetcode 1402 Solution

This article provides solution to leetcode question 1402 (count-square-submatrices-with-all-ones)

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

Solution

import numpy as np

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

        dp = np.ndarray((m, n)).tolist()

        dp[0][0] = matrix[0][0]

        for i in range(m):
            dp[i][0] = matrix[i][0]

        for j in range(n):
            dp[0][j] = matrix[0][j]

        for i in range(1, m):
            for j in range(1, n):
                if matrix[i][j] == 0:
                    dp[i][j] = 0
                    continue

                max_size = min(dp[i - 1][j], dp[i][j - 1])

                if matrix[i - max_size][j - max_size] == 1:
                    dp[i][j] = max_size + 1
                else:
                    dp[i][j] = max_size

        ans = 0
        for i in range(m):
            for j in range(n):
                ans += dp[i][j]
        return ans