二维前缀和

约 976 字大约 3 分钟

2026-01-20

来自1292. 元素和小于等于阈值的正方形的最大边长，利用容斥原理思想构造二维前缀和数组，进行遍历。

前缀和可以理解为对数组进行积分（只是抽象理解，这样的说法完全不科学）。假设有长度为 $n$ 数组 $\{ a_i \}$ ，前缀和数组 $\{ S_i \}$ 定义为 $S_i = \sum_{k=0}^{i} a_k$ 。当然，有时候根据题目的意思，也可以写作 $S_i = \sum_{k=0}^{i-1} a_k$ （当 $i = 0$ 时， $S_i = 0$ ）。前缀和与差分数组一样，极大地方便了我们对区间进行查询和修改，避免了大量的重复计算。当我们想要询问区间 $[l, r]$ 的和时，可以通过前缀和数组 $S$ 快速计算出来： $\sum_{i=l}^{r} a_i = S_{r} - S_{l-1}$ ，此时通过 $O(n)$ 的时间预处理前缀和数组，可以让每次查询的时间复杂度降为 $O(1)$ 。

C++ 标准库中实现了前缀和函数 std::partial_sum ，在头文件 <numeric> 中定义，可以直接使用：

#include <numeric>
#include <vector>

std::vector<int> a = {1, 2, 3, 4, 5};
std::vector<int> prefix_sum(a.size() + 1);
std::partial_sum(a.begin(), a.end(), prefix_sum.begin() + 1); // prefix_sum = {0, 1, 3, 6, 10, 15}

// or you can
prefix_sum.resize(a.size());
std::partial_sum(a.begin(), a.end(), prefix_sum.begin(), std::plus<int>()); // prefix_sum = {1, 3, 6, 10, 15}

二维前缀和

有了前缀和的概念，我们可以将其推广到二维空间。假设有一个 $m \times n$ 的矩阵 $\{ a_{i,j} \}$ ，我们定义二维前缀和矩阵 $\{ S_{i,j} \}$ 为：

S_{i,j} = \sum_{p=0}^{i} \sum_{q=0}^{j} a_{p,q}

利用容斥原理避免重复计算，我们有

S_{i,j} = a_{i,j} + S_{i-1,j} + S_{i,j-1} - S_{i-1,j-1}

如果我们要查询矩阵中某个子矩阵 $\{ a_{i,j} \}$ （左上角为 $(x_1, y_1)$ ，右下角为 $(x_2, y_2)$ ）的元素和，可以通过二维前缀和矩阵快速计算出来：

\sum_{i=x_1}^{x_2} \sum_{j=y_1}^{y_2} a_{i,j} = S_{x_2,y_2} - S_{x_1-1,y_2} - S_{x_2,y_1-1} + S_{x_1-1,y_1-1}

利用 $O(mn)$ 的时间预处理二维前缀和矩阵，可以让每次查询子矩阵和的时间复杂度降为 $O(1)$ 。

解答

由于题目给的数据不大，我们采用二分方法，每次判断时暴力枚举所有符合边长的正方形，利用二维前缀和对每个区域内的和进行查询。

class Solution {
public:
    int maxSideLength(vector<vector<int>>& mat, int threshold) {
        int m = mat.size();
        int n = mat[0].size();
        vector<vector<int>> presum(m +1, vector<int>(n + 1, 0));
        for (int i = 1; i <= m; i++) {
            for (int j = 1; j <= n; j++) {
                presum[i][j] = presum[i - 1][j] + presum[i][j - 1] - presum[i - 1][j - 1] + mat[i - 1][j - 1];
            }
        }

        auto check = [&](int mid) -> bool {
            int t = 0;
            for (int i = 1; i <= m -  mid + 1; i++) {
                for (int j = 1; j <= n - mid + 1; j++) {
                    if (presum[i + mid - 1][j + mid - 1] - presum[i - 1][j + mid - 1] - presum[i + mid - 1][j - 1] + presum[i - 1][j - 1] <= threshold) {
                        return true;
                    }
                }
            }
            return false;
        };
        int l = 0, r = min(m, n);
        int ans = 0;

        while (l <= r) {
            int mid = l + (r - l) / 2;
            if (check(mid)) {
                ans = mid;
                l = mid + 1;
            } else {
                r = mid - 1;
            }
        }
        return ans;
    }
};

二分法分析

二分寻找答案分两种情况。一种是lower_bound，另一种是upper_bound。此处upper_bound不同于STL中的upper_bound，STL的upper_bound是寻找第一个大于某个值的位置。

lower_bound：寻找第一个满足条件的值。即找到最小的满足条件的值。

while (l < r) {
    int mid = l + (r - l) / 2;
    if (check(mid)) {
        r = mid; // 注意这里不是 mid - 1
    } else {
        l = mid + 1;
    }
}

upper_bound：寻找最后一个满足条件的值。即找到最大的满足条件的值。

while (l <= r) {
    int mid = l + (r - l) / 2;
    if (check(mid)) {
        ans = mid;
        l = mid + 1;
    } else {
        r = mid - 1;
    }
}

常规我们对一个数组进行二分查找精确值时，一般使用lower_bound写法。

二维前缀和

前置知识

前缀和

二维前缀和

解答

二分法分析