题目描述

给你一个 m * n 的矩阵，矩阵中的元素不是 0 就是 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

输入：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

(暴力枚举) $O(nm \min(n, m))$

1. 预处理二维的前缀和数组 sum。
2. 枚举正方形的边长 s，从 1 到 $\min(n, m)$。
3. 枚举每一个合法的右下家的位置 (i, j)，如果在以 (i, j) 为右下角的正方形总和为 s，则答案加 1。

时间复杂度

• 预处理的时间复杂度为 $O(nm)$。
• 共有 $\min(n, m)$ 种正方形，故总时间复杂度为 $O(nm \min(n, m))$。

空间复杂度

• 需要额外 $O(nm)$ 的空间存储前缀和数组。

C++ 代码

class Solution {
public:
    int countSquares(vector<vector<int>>& matrix) {
        int n = matrix.size(), m = matrix[0].size();
        vector<vector<int>> sum(n + 1, vector<int>(m + 1, 0));

        for (int i = 1; i <= n; i++)
            for (int j = 1; j <= m; j++)
                sum[i][j] = sum[i - 1][j] + sum[i][j - 1]
                            - sum[i - 1][j - 1] + matrix[i - 1][j - 1];

        int ans = 0;

        for (int s = 1; s <= min(n, m); s++)
            for (int i = s; i <= n; i++)
                for (int j = s; j <= m; j++)
                    if (sum[i][j] - sum[i - s][j]
                        - sum[i][j - s] + sum[i - s][j - s] == s * s)
                        ans++;

        return ans;
    }
};

算法2

(动态规划) $O(nm)$

1. 设 $f(i, j)$ 是以 (i, j) 为右下角的最大正方形的边长。则以 (i, j) 为右下角的正方形的个数就是 $f(i, j)$。
2. 如果 matrix[i][j] == 1，则 $f(i, j) = \min(f(i - 1, j), f(i, j - 1), f(i - 1, j - 1)) + 1$；否则 $f(i, j) = 0$。
3. 最终答案就是 $\sum f(i, j)$，即所有位置的 $f(i, j)$ 求和。

时间复杂度

• 状态数为 $O(nm)$ 个，每次转移时间为常数，故总时间复杂度为 $O(nm)$。

空间复杂度

• 需要额外 $O(nm)$ 的空间存储状态，但可以通过滚动数组将空间优化到 $O(n)$。
• 如果允许，也可以使用原数组的空间，从而空间复杂度降为 $O(1)$。

C++ 代码

class Solution {
public:
    int countSquares(vector<vector<int>>& matrix) {
        int n = matrix.size(), m = matrix[0].size();
        vector<vector<int>> f(n + 1, vector<int>(m + 1, 0));

        int ans = 0;
        for (int i = 1; i <= n; i++)
            for (int j = 1; j <= m; j++)
                if (matrix[i - 1][j - 1] == 1) {
                    f[i][j] = min(min(f[i - 1][j], f[i][j - 1]), f[i - 1][j - 1]) + 1;
                    ans += f[i][j];
                }

        return ans;
    }
};