2D Prefix Sum

MrHe··3 min read

二维前缀和

1. 问题场景：从一维到二维

我们先回忆一下一维前缀和。给你一个数组 arr，如果频繁地问你 arr[L...R] 的和，我们会预处理一个前缀和数组 preSum，其中 preSum[i] = arr[0] + ... + arr[i]。这样 sum(L, R) 就可以用 preSum[R] - preSum[L-1] 在 O(1) 时间内得到。

现在，我们把问题升级到二维： 给你一个二维矩阵 matrix，现在会非常频繁地问你：以 (r1, c1) 为左上角，以 (r2, c2) 为右下角的子矩阵，所有元素的和是多少？

暴力解法： 最直观的想法就是，每来一个查询，我就写一个二重循环去遍历这个子矩阵，然后把所有元素累加起来。

// 暴力解法
int sum = 0;
for (int i = r1; i <= r2; i++) {
    for (int j = c1; j <= c2; j++) {
        sum += matrix[i][j];
    }
}

复杂度分析：如果矩阵大小为 N x M，一个子矩阵最大也可以是 N x M。那么单次查询的复杂度最坏是 O(N M)。如果有 K 次查询，总复杂度就是 O(K N * M)。在 K、N、M 都很大的情况下，这个方法肯定会超时。

思路演进： 既然一维问题能用前缀和优化到 O(1) 查询，那二维问题也一定有类似的技巧！我们能不能也预处理一个“前缀和矩阵”，让查询也变成 O(1) 呢？

答案是肯定的，这个“前缀和矩阵”就是我们的主角——二维前缀和数组。

2. 二维前缀和的原理

核心定义： 我们定义一个二维前缀和数组，叫做 preSum。preSum[i][j] 的含义是：从矩阵左上角 (0, 0) 到当前位置 (i, j) 所构成的矩形区域内，所有元素的和。

如何构建 preSum 数组？ 计算 preSum[i][j] 时，我们需要利用已经计算好的、更小的前缀和。这里就要用到一个非常重要的数学思想——容斥原理 (Inclusion-Exclusion Principle)。

(图片来源: halfrost/LeetCode-Go)

看上图，为了求 preSum[i][j] (即右下角为 (i,j) 的整个大矩形)，我们可以这样做：

  1. 先拿它上面的那个矩形的和，即 preSum[i-1][j] (图中区域 A + B)。

  2. 再拿它左边的那个矩形的和，即 preSum[i][j-1] (图中区域 A + C)。

  3. 这时你会发现，左上角那个叫 A 的区域 preSum[i-1][j-1] 被加了两次，所以要减掉一次。

  4. 最后，我们把这三块区域加加减减，得到的是 preSum[i][j] 这个大矩形中除了 matrix[i][j] 本身（图中区域 D）以外所有元素的和。所以，最后必须把 matrix[i][j] 加上。

于是，我们得到了构建二维前缀和的递推公式： preSum[i][j] = preSum[i-1][j] + preSum[i][j-1] - preSum[i-1][j-1] + matrix[i][j]

这个构建过程的时间复杂度是 O(N * M)。

如何使用 preSum 数组进行查询（区域和）？ 我们的目标是求 (r1, c1) 到 (r2, c2) 子矩阵的和。同样，还是用容斥原理！

看上图，我们想要求绿色区域的和：

  1. 首先，我们取出从 (0,0) 到右下角 (r2, c2) 的大矩形的和，即 preSum[r2][c2]。

  2. 这个大矩形包含了我们不想要的部分。我们先减掉它上方多余的矩形，即从(0,0)到(r1-1, c2)的和 preSum[r1-1][c2]。

  3. 再减掉它左方多余的矩形，即从(0,0)到(r2, c1-1)的和 preSum[r2][c1-1]。

  4. 此时你会发现，左上角那个从 (0,0) 到 (r1-1, c1-1) 的小矩形，被减了两次！所以，我们要把它加回来一次，即 + preSum[r1-1][c1-1]。

于是，我们得到了查询公式： sum(r1, c1, r2, c2) = preSum[r2][c2] - preSum[r1-1][c2] - preSum[r2][c1-1] + preSum[r1-1][c1-1]

每一次查询，都只是四个点的数组访问和三次加减法，时间复杂度是 O(1)！

3. OI/ACM 风格代码实现

这道题就是 LeetCode 的原题。

问题重述 实现 NumMatrix 类：

  • NumMatrix(int[][] matrix) 初始化对象。

  • int sumRegion(int r1, int c1, int r2, int c2) 返回子矩阵的和。

实现技巧：1-based 索引 在写代码时，你会发现 i-1, j-1, r1-1, c1-1 这些操作很容易在 i, j, r1, c1 为 0 时导致数组越界。为了避免写一堆 if-else 来处理边界，竞赛中有一个非常优雅的常用技巧：将前缀和数组 preSum 的尺寸定义为 (N+1) x (M+1)。

这样，我们的 preSum 数组的有效索引就从 (1, 1) 开始，preSum[i][j] 对应原矩阵 matrix[i-1][j-1]。而 preSum 的第 0 行和第 0 列天然就是 0，这使得我们的递推公式和查询公式无需任何特殊处理，可以完美地统一起来。

/**
 * LeetCode 304. 二维区域和检索 - 矩阵不可变
 * OI/ACM 风格实现，采用 1-based 索引技巧
 */
class NumMatrix {
​
    // preSum[i][j] 存储的是 matrix 中 (0,0) 到 (i-1,j-1) 的矩形和
    private int[][] preSum;
​
    /**
     * 构造函数：预处理二维前缀和数组
     * 时间复杂度: O(rows * cols)
     * @param matrix 原始矩阵
     */
    public NumMatrix(int[][] matrix) {
        if (matrix == null || matrix.length == 0 || matrix[0].length == 0) {
            // 处理空矩阵的边界情况
            return;
        }
        int rows = matrix.length;
        int cols = matrix[0].length;
​
        // 核心技巧：创建 (rows + 1) x (cols + 1) 的前缀和数组
        // 这样 preSum 的第 0 行和第 0 列都为 0，完美处理边界
        preSum = new int[rows + 1][cols + 1];
​
        // 遍历原矩阵，构建前缀和数组
        for (int i = 1; i <= rows; i++) {
            for (int j = 1; j <= cols; j++) {
                // 应用递推公式
                // preSum 的 (i, j) 对应 matrix 的 (i-1, j-1)
                preSum[i][j] = preSum[i - 1][j] + preSum[i][j - 1] - preSum[i - 1][j - 1] + matrix[i - 1][j - 1];
            }
        }
    }
​
    /**
     * 查询子矩阵的和
     * 时间复杂度: O(1)
     * @param row1, col1, row2, col2  (0-based 索引)
     * @return 子矩阵的和
     */
    public int sumRegion(int row1, int col1, int row2, int col2) {
        if (preSum == null) {
            return 0;
        }
​
        // 将输入的 0-based 坐标转换为我们 1-based 的 preSum 数组坐标
        int r1 = row1 + 1;
        int c1 = col1 + 1;
        int r2 = row2 + 1;
        int c2 = col2 + 1;
​
        // 应用查询公式
        return preSum[r2][c2] - preSum[r1 - 1][c2] - preSum[r2][c1 - 1] + preSum[r1 - 1][c1 - 1];
    }
}
  • #5ymn57ya5zkm