Skip to main content

Command Palette

Search for a command to run...

2D Prefix Sum

Published
3 min readView as Markdown

二维前缀和


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. 二维前缀和的原理

核心定义: 我们定义一个二维前缀和数组,叫做 preSumpreSum[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];
    }
}

More from this blog

字典树和后缀数组

14.2 字典树的应用算法设计 字典树,也叫Trie树或者前缀树,顾名思义,它就是一种专门用来处理字符串前缀的树形结构。每个节点代表一个字符,从根节点到任意一个节点的路径,就构成了一个字符串。它的核心思想就是用空间换时间,利用字符串的公共前缀来降低查询时间的开销。 14.2.1 实现Trie(前缀树) 问题描述 实现一个 Trie (前缀树),包含 insert, search, 和 startsWith 这三个操作。 Trie() 初始化前缀树对象。 void insert(String...

Oct 4, 20259 min read

树状数组( Binary Indexed Tree)

树状数组(Fenwick Tree / Binary Indexed Tree)这个东西。这玩意儿听起来挺唬人,但本质上就是个用数组模拟树形结构的工具,专门用来解决一类问题:频繁的单点更新和区间查询。如果我们用普通数组,更新是 O(1),但查区间和就是 O(N);用前缀和数组,查区间和是 O(1),但你更新一个点,后面所有的前缀和都得改,又是 O(N)。两边都不讨好。树状数组就把这两个操作都干到了 O(logN),一下就平衡了。 废话不多说,咱们从题入手,看看这东西到底怎么用,思路是怎么一步步升...

Oct 4, 202510 min read

线段树-树形结构

线段树,顾名思义,就是用来处理区间或线段问题的树形结构。它能高效地解决一类问题:对一个数组进行频繁的区间查询和元素更新。如果你看到一个问题,反复地问你“从L到R这个范围内的xx是什么?”或者“把L到R这个范围内的数都改成xx”,那线札树可能就是你的答案。 它的核心思想就是“分治”。把一个大区间,一分为二,变成两个小区间,直到每个区间只包含一个元素。这样,一个长度为N的数组,就对应了一棵满二叉树(或近似满二叉树),树的深度是 O(logN)。每一次查询和更新,我们都只需要沿着树的一条路径走到底,所...

Oct 4, 20259 min read

优先队列PriorityQueue

很多时候,我们处理数据不只是简单地存进去、取出来,而是希望每次取出的都是当前这批数据里"最"牛的那个,比如最大、最小、最重要等等。普通队列是先进先出,栈是后进先出,都满足不了这个需求。这时候,优先队列就登场了。 你可以把它想象成一个VIP候机室,不管你什么时候来的,只要你的"优先级"最高(比如是头等舱、白金会员),你就能最先登机。在算法里,这个"优先级"通常就是元素的大小。 Java里 PriorityQueue 就是它的标准实现,底层是一棵二叉堆(小顶堆)。这意味着,它能用 O(logN) 的...

Oct 4, 202514 min read

并查集(Disjoint Set Union, DSU)

并查集(Disjoint Set Union, DSU)是一种非常有用的数据结构,主要用来处理一些不相交集合的合并及查询问题。听起来很抽象,但其实核心就两个操作: find:查找一个元素属于哪个集合(通常通过返回该集合的代表元素,也就是根节点)。 union:将两个元素所在的集合合并成一个集合。 为了让find操作更快,我们通常会加上一个“路径压缩”的优化。为了让合并后的树结构更平衡,会加上“按秩合并”或“按大小合并”的优化。 10.2 一维并查集的应用算法设计 10.2.1 以图判树...

Oct 4, 202512 min read

小何的博客

49 posts

欢迎来到小何的博客,25届本科毕业生,河南科技大学软件工程专业,主攻Java,对python、vue、uniapp、大模型等技术都有涉及。