LeetCode 1277. Count Square Submatrices with All Ones

LeetCode 1277

第一个想到的方法是用dp，把每行，每列的1的长度求出来，然后再算当前位置square的长度。

        m = len(matrix)
        n = len(matrix[0])
        #dp1,dp2, 表示当前位置，行列1的长度
        dp1 = [[0]*n for i in range(0,m)]
        dp2 = [[0]*n for i in range(0,m)]
        #dp3表示当前位置square长度
        dp3 = [[0]*n for i in range(0,m)]
        count = 0

        for i in range(0, m):
            for j in range(0, n):
                if i == 0 and j == 0:
                    dp1[i][j] = dp2[i][j] = dp3[i][j] = matrix[0][0]
                    count += dp3[i][j]
                    continue
                if i == 0:
                    dp1[i][j] = dp1[i][j-1] + 1 if matrix[i][j] == 1 else 0
                    dp2[i][j] = dp3[i][j] = matrix[i][j]
                    count += dp3[i][j]
                    continue
                if j == 0:
                    dp2[i][j] = dp1[i-1][j] + 1 if matrix[i][j] == 1 else 0
                    dp1[i][j] = dp3[i][j] = matrix[i][j]
                    count += dp3[i][j]
                    continue
                dp1[i][j] = dp1[i][j-1] + 1 if matrix[i][j] == 1 else 0
                dp2[i][j] = dp2[i-1][j] + 1 if matrix[i][j] == 1 else 0
                if matrix[i][j] == 1:
                    dp3[i][j] = min(dp1[i][j], dp2[i][j], dp3[i-1][j-1] + 1)
                    count += dp3[i][j]

        return count

有了这个之后，其实我们发现如果允许改变当前matrix的值的话，我们可以简化掉：

    def countSquares(self, matrix: List[List[int]]) -> int:
        if not matrix:
            return 0
        
        rows = len(matrix)
        cols = len(matrix[0])
        
        ans = 0
        for r in range(rows):
            for c in range(cols):
                if matrix[r][c] == 1:
                    if r>0 and c>0:
                        matrix[r][c] = min(matrix[r][c-1], matrix[r-1][c], matrix[r-1][c-1])+1
                    ans += matrix[r][c]
        return ans