Leetcode 1277：统计全为 1 的正方形子矩阵（超详细的解法！！！）

给你一个 m * n 的矩阵，矩阵中的元素不是 0 就是 1，请你统计并返回其中完全由 1 组成的 正方形 子矩阵的个数。

示例 1：

输入：matrix =
[
  [0,1,1,1],
  [1,1,1,1],
  [0,1,1,1]
]
输出：15
解释： 
边长为 1 的正方形有 10 个。
边长为 2 的正方形有 4 个。
边长为 3 的正方形有 1 个。
正方形的总数 = 10 + 4 + 1 = 15.

示例 2：

输入：matrix = 
[
  [1,0,1],
  [1,1,0],
  [1,1,0]
]
输出：7
解释：
边长为 1 的正方形有 6 个。 
边长为 2 的正方形有 1 个。
正方形的总数 = 6 + 1 = 7.

提示：

• 1 <= arr.length <= 300
• 1 <= arr[0].length <= 300
• 0 <= arr[i][j] <= 1

解题思路

首先可以想到暴力解法，也就是先枚举所有的正方形，然后判断其中的数字是不是都是1。但是这么做的话，时间复杂度太高了。我们可以先通过前缀和的方式统计好当前位置(i,j)为右下角，(0,0)为左上角的矩形中1的个数。这样我们最后只需要枚举所有的正方形即可，时间复杂度就是O(n^3)。

class Solution:
    def countSquares(self, data: List[List[int]]) -> int:
        r, c = len(data), len(data[0])
        arr, res = [[0] * (c + 1) for _ in range(r + 1)], 0

        for i in range(1, r + 1):
            for j in range(1, c + 1):
                if data[i-1][j-1]:
                    res += 1
                    arr[i][j] = 1
                arr[i][j] += arr[i-1][j] + arr[i][j-1] - arr[i-1][j-1]

        for l in range(2, min(r, c) + 1):
            for i in range(r - l + 1):
                for j in range(c - l + 1):
                    if arr[i+l][j+l] + arr[i][j] - arr[i][j+l]-arr[i+l][j] == l**2:
                        res += 1
        return res

这个问题还可以参考之前Leetcode 221：最大正方形中的思想，使用动态规划来解。定义函数$f(i,j)$表示右下角坐标为(i,j)的最大正方形边长，那么

• $f(i,j)=min(f(i-1,j),f(i,j-1),f(i-1,j-1))+1$

知道最大正方形边长有什么用呢？实际上最大正方形边长就是以(i,j)为右下角的正方形个数。我们只需要将所有边长加起来即可。

class Solution:
    def countSquares(self, A: List[List[int]]) -> int:
        res = 0
        for i in range(len(A)):
            for j in range(len(A[0])):
                if i and j and A[i][j]:
                    A[i][j] = min(A[i - 1][j], A[i][j - 1], A[i - 1][j - 1]) + 1
                res += A[i][j]
        return res

实际上就是之前问题中的res=max(res,)换成了res+=。

我将该问题的其他语言版本添加到了我的GitHub Leetcode

如有问题，希望大家指出！！！