【算法思考记录】【前缀和，C++】力扣1277. 统计全为 1 的正方形子矩阵

原题链接

---

使用前缀和算法解决统计全为1的正方形子矩阵问题

题目分析

题目要求我们统计在一个由0和1构成的矩阵中，所有完全由1组成的正方形子矩阵的数量。这是一道中等难度的算法题目，其关键在于高效地计算出不同大小的正方形子矩阵是否完全由1组成。

解题思路

解决此问题的一个有效方法是使用前缀和算法。前缀和是一种预处理技术，通过计算数组中每个元素对应的前缀和，可以快速计算出任意子数组的和。在这个问题中，我们将前缀和算法扩展到二维，以便快速计算任意子矩阵的元素和。

前缀和算法的基本原理

前缀和算法是一种用于快速求解数组或矩阵中某个区间内元素和的技术。这种方法特别适用于需要多次对同一数据集的不同区间求和的情况。它通过预处理数据来优化求和的效率。

一维前缀和

假设我们有一个一维数组 arr，其前缀和数组 preSum 可以这样计算：

1. 初始化 preSum[0] = 0。
2. 对于 arr 中的每个元素 arr[i]，计算 preSum[i+1] = preSum[i] + arr[i]。

这样，preSum[i] 存储了 arr 中从第一个元素到第 i-1 个元素的总和。如果我们想要计算 arr 中从第 l 个元素到第 r 个元素的和，我们只需要计算 preSum[r+1] - preSum[l]。

二维前缀和

二维前缀和是一维前缀和的扩展，用于处理矩阵。假设我们有一个二维数组 matrix，其对应的二维前缀和矩阵 preSum 可以这样计算：

1. 初始化 preSum 的所有元素为0。
2. 对于 matrix 中的每个元素 matrix[i][j]，计算 preSum[i+1][j+1] 的值。这个值等于其左侧元素的前缀和、上方元素的前缀和和左上方元素的前缀和的总和，再减去左上方元素的前缀和，最后加上 matrix[i][j] 本身的值。

计算矩阵中某个子矩形区域内元素的总和时，设该区域的左上角坐标为 (r1, c1)，右下角坐标为 (r2, c2)，则该区域内元素的总和可以表示为：

preSum[r2][c2] - preSum[r1-1][c2] - preSum[r2][c1-1] + preSum[r1-1][c1-1]

这个公式的含义是：从整个矩阵的起点到 (r2, c2) 的总和，减去上方和左方的部分重叠的区域，再加上左上角重复减去的部分。

应用

在解决我们的问题时，我们利用了二维前缀和算法来快速计算矩阵中任意正方形区域内的元素总和。通过这种方式，我们避免了对每个可能的正方形区域进行繁琐的重复计算，显著提高了算法的效率。

代码实现

以下是基于前缀和算法的解决方案，包含详细的注释说明。

#include 
using namespace std;

class Solution {
public:
    int countSquares(vector<vector<int>>& matrix) {
        int m = matrix.size(), n = matrix[0].size();
        // 使用二维数组pre_sum存储前缀和
        vector<vector<int>> pre_sum(m + 1, vector<int>(n + 1));

        // 构建前缀和矩阵
        for (int i = 0; i < m; ++i) {
            for (int j = 0; j < n; ++j) {
                pre_sum[i + 1][j + 1] = pre_sum[i][j + 1] + pre_sum[i + 1][j] - pre_sum[i][j] + matrix[i][j];
            }
        }

        int ans = 0;
        // max_edge为可能的最大正方形边长
        int max_edge = min(m, n);

        // 遍历所有可能的正方形大小和位置
        for (int k = 1; k <= max_edge; ++k) {
            for (int i = k; i <= m; ++i) {
                for (int j = k; j <= n; ++j) {
                    // 计算正方形的左上角和右下角坐标
                    int r1 = i - k, c1 = j - k;
                    int r2 = i, c2 = j;
                    // 计算正方形内部的元素和
                    long long actual = 0LL + pre_sum[r2][c2] - pre_sum[r1][c2] - pre_sum[r2][c1] + pre_sum[r1][c1];
                    // 如果和等于正方形面积，则说明该正方形完全由1组成
                    if (actual == k * k) {
                        ans++;
                    }
                }
            }
        }
        return ans;
    }
};

算法解析

1. 初始化前缀和矩阵：我们创建一个二维数组pre_sum来存储原矩阵的前缀和。pre_sum[i + 1][j + 1]表示原矩阵左上角到(i, j)位置形成的矩形区域内所有元素的和。

2. 计算前缀和：通过遍历原矩阵，利用pre_sum[i][j] = pre_sum[i-1][j] + pre_sum[i][j-1] - pre_sum[i-1][j-1] + matrix[i-1][j-1]公式计算每个位置的前缀和。

3. 遍历所有可能的正方形：我们遍历所有可能的正方形大小和起始位置，利用前缀和快速计算正方形内部的元素和。

4. 统计符合条件的正方形：对于每个正方形，如果其内部元素和等于其面积（即边长的平方），则说明这个正方形完全由1组成，计数器ans增加。

5. 返回结果：最后返回计数器

ans，即为完全由1组成的正方形子矩阵的总数。

结论

通过使用前缀和算法，我们能够有效地解决这个问题，减少了重复计算，提高了算法的效率。这种方法在处理涉及子数组或子矩阵和的问题时非常有用。