小而美的算法技巧：前缀和数组

本文讲解的例题

LeetCode	力扣	难度
303. Range Sum Query - Immutable	303. 区域和检索 - 数组不可变	🟢
304. Range Sum Query 2D - Immutable	304. 二维区域和检索 - 矩阵不可变	🟠

前置知识

阅读本文前，你需要先学习：

数组基础

前缀和技巧适用于快速、频繁地计算一个索引区间内的元素之和。

一维数组中的前缀和

先看一道例题，力扣第 303 题「区域和检索 - 数组不可变」，让你计算数组区间内元素的和，这是一道标准的前缀和问题：

303. 区域和检索 - 数组不可变 | 力扣 | LeetCode | 🟢

给定一个整数数组 nums，处理以下类型的多个查询:

计算索引 left 和 right （包含 left 和 right）之间的 nums 元素的和，其中 left <= right

实现 NumArray 类：

NumArray(int[] nums) 使用数组 nums 初始化对象
int sumRange(int i, int j) 返回数组 nums 中索引 left 和 right 之间的元素的总和，包含 left 和 right 两点（也就是 nums[left] + nums[left + 1] + ... + nums[right] )

示例 1：

输入：
["NumArray", "sumRange", "sumRange", "sumRange"]
[[[-2, 0, 3, -5, 2, -1]], [0, 2], [2, 5], [0, 5]]
输出：
[null, 1, -1, -3]

解释：
NumArray numArray = new NumArray([-2, 0, 3, -5, 2, -1]);
numArray.sumRange(0, 2); // return 1 ((-2) + 0 + 3)
numArray.sumRange(2, 5); // return -1 (3 + (-5) + 2 + (-1)) 
numArray.sumRange(0, 5); // return -3 ((-2) + 0 + 3 + (-5) + 2 + (-1))

提示：

1 <= nums.length <= 10⁴
-10⁵ <= nums[i] <= 10⁵
0 <= i <= j < nums.length
最多调用 10⁴ 次 sumRange 方法

题目来源：力扣 303. 区域和检索 - 数组不可变。

java

// 题目要求你实现这样一个类
class NumArray {

    public NumArray(int[] nums) {}
    
    // 查询闭区间 [left, right] 的累加和
    public int sumRange(int left, int right) {}
}

cpp

python

javascript

sumRange 函数需要计算并返回一个索引区间之内的元素和，没学过前缀和的人可能写出如下代码：

class NumArray {
	private int[] nums;
	public NumArray(int[] nums) {
		this.nums = nums;
	}

	public int sumRange(int left, int right) {
		// 用 for 循环遍历求和
		int sum = 0;
		for (int i = left; i <= right; i++) {
			sum += nums[i];
		}
		return sum;
	}
}

这个解法每次调用 sumRange 函数时，都要进行一次 for 循环遍历，时间复杂度为 $O(N)$ ，而 sumRange 的调用频率可能非常高，所以这个算法的效率很低。

正确的解法是使用前缀和技巧进行优化，使得 sumRange 函数的时间复杂度为 $O(1)$ ：

java

class NumArray {
    // 前缀和数组
    private int[] preSum;

    // 输入一个数组，构造前缀和
    public NumArray(int[] nums) {
        // preSum[0] = 0，便于计算累加和
        preSum = new int[nums.length + 1];
        // 计算 nums 的累加和
        for (int i = 1; i < preSum.length; i++) {
            preSum[i] = preSum[i - 1] + nums[i - 1];
        }
    }

    // 查询闭区间 [left, right] 的累加和
    public int sumRange(int left, int right) {
        return preSum[right + 1] - preSum[left];
    }
}

cpp

python

javascript

核心思路是我们 new 一个新的数组 preSum 出来，preSum[i] 记录 nums[0..i-1] 的累加和，看图 $10 = 3 + 5 + 2$ ：

看这个 preSum 数组，如果我想求索引区间 [1, 4] 内的所有元素之和，就可以通过 preSum[5] - preSum[1] 得出。

这样，sumRange 函数仅仅需要做一次减法运算，避免了每次进行 for 循环调用，最坏时间复杂度为常数 $O(1)$ 。

你可以点开下面的可视化动画，点击 preSum[i] = preSum[i - 1] + nums[i - 1] 这行代码，即可看到 preSum 数组的计算，多次点击 console.log 这行代码，即可看到 sumRange 函数的调用：

算法可视化面板

这个技巧在生活中运用也挺广泛的，比方说，你们班上有若干同学，每个同学有一个期末考试的成绩（满分 100 分），那么请你实现一个 API，输入任意一个分数段，返回有多少同学的成绩在这个分数段内。

那么，你可以先通过计数排序的方式计算每个分数具体有多少个同学，然后利用前缀和技巧来实现分数段查询的 API：

// 存储着所有同学的分数
int[] scores = new int[]{...};
// 试卷满分 100 分
int[] count = new int[100 + 1];

// 记录每个分数有几个同学
for (int score : scores) {
    count[score]++;
}
// 构造前缀和数组
for (int i = 1; i < count.length; i++) {
    count[i] = count[i] + count[i-1];
}

// 利用 count 这个前缀和数组进行分数段查询

// 查询分数在 [80, 90] 之间的同学有多少人
int result = count[90] - count[79];

接下来，我们看一看前缀和思路在二维数组中如何运用。

二维矩阵中的前缀和

这是力扣第 304 题「二维区域和检索 - 矩阵不可变」，其实和上一题类似，上一题是让你计算子数组的元素之和，这道题让你计算二维矩阵中子矩阵的元素之和：

304. 二维区域和检索 - 矩阵不可变 | 力扣 | LeetCode | 🟠

给定一个二维矩阵 matrix，以下类型的多个请求：

计算其子矩形范围内元素的总和，该子矩阵的 左上角 为 (row1, col1) ，右下角 为 (row2, col2) 。

实现 NumMatrix 类：

NumMatrix(int[][] matrix) 给定整数矩阵 matrix 进行初始化
int sumRegion(int row1, int col1, int row2, int col2) 返回 左上角 (row1, col1) 、右下角 (row2, col2) 所描述的子矩阵的元素总和。

示例 1：

输入: 
["NumMatrix","sumRegion","sumRegion","sumRegion"]
[[[[3,0,1,4,2],[5,6,3,2,1],[1,2,0,1,5],[4,1,0,1,7],[1,0,3,0,5]]],[2,1,4,3],[1,1,2,2],[1,2,2,4]]
输出: 
[null, 8, 11, 12]

解释:
NumMatrix numMatrix = new NumMatrix([[3,0,1,4,2],[5,6,3,2,1],[1,2,0,1,5],[4,1,0,1,7],[1,0,3,0,5]]);
numMatrix.sumRegion(2, 1, 4, 3); // return 8 (红色矩形框的元素总和)
numMatrix.sumRegion(1, 1, 2, 2); // return 11 (绿色矩形框的元素总和)
numMatrix.sumRegion(1, 2, 2, 4); // return 12 (蓝色矩形框的元素总和)

提示：

m == matrix.length
n == matrix[i].length
1 <= m, n <= 200
-10⁵ <= matrix[i][j] <= 10⁵
0 <= row1 <= row2 < m
0 <= col1 <= col2 < n
最多调用 10⁴ 次 sumRegion 方法

题目来源：力扣 304. 二维区域和检索 - 矩阵不可变。

当然，你可以用一个嵌套 for 循环去遍历这个矩阵，但这样的话 sumRegion 函数的时间复杂度就高了，你算法的格局就低了。

注意任意子矩阵的元素和可以转化成它周边几个大矩阵的元素和的运算：

而这四个大矩阵有一个共同的特点，就是左上角都是 (0, 0) 原点。

那么做这道题更好的思路和一维数组中的前缀和是非常类似的，我们可以维护一个二维 preSum 数组，专门记录以原点为顶点的矩阵的元素之和，就可以用几次加减运算算出任何一个子矩阵的元素和：

java

class NumMatrix {
    // preSum[i][j] 记录矩阵 [0, 0, i-1, j-1] 的元素和
    private int[][] preSum;

    public NumMatrix(int[][] matrix) {
        int m = matrix.length, n = matrix[0].length;
        if (m == 0 || n == 0) return;
        // 构造前缀和矩阵
        preSum = new int[m + 1][n + 1];
        for (int i = 1; i <= m; i++) {
            for (int j = 1; j <= n; j++) {
                // 计算每个矩阵 [0, 0, i, j] 的元素和
                preSum[i][j] = preSum[i-1][j] + preSum[i][j-1] + matrix[i - 1][j - 1] - preSum[i-1][j-1];
            }
        }
    }

    // 计算子矩阵 [x1, y1, x2, y2] 的元素和
    public int sumRegion(int x1, int y1, int x2, int y2) {
        // 目标矩阵之和由四个相邻矩阵运算获得
        return preSum[x2+1][y2+1] - preSum[x1][y2+1] - preSum[x2+1][y1] + preSum[x1][y1];
    }
}

cpp

python

javascript

这样，sumRegion 函数的时间复杂度也用前缀和技巧优化到了 $O(1)$ ，这是典型的「空间换时间」思路。

你可以点开下面的可视化动画，多次点击 preSum[i][j] = ... 这一行代码，即可看到 preSum 数组的计算过程，多次点击 console.log 这一行代码，即可看到 sumRegion 函数的调用：

算法可视化面板

前缀和技巧就讲到这里，应该说这个算法技巧是会者不难难者不会，实际运用中还是要多培养自己的思维灵活性，做到一眼看出题目是一个前缀和问题。

拓展延伸

本文讲解的前缀和技巧是利用预计算的 preSum 数组快速计算索引区间内的元素和，但实际上它不仅仅局限于求和，也可以用来快速计算区间乘积等其他场景。

而且前缀和数组经常和其他数据结构或算法技巧相结合，我会在前缀和技巧高频习题中结合习题讲解。

但是，前缀和技巧有几个局限性。

第一个局限性：使用前缀和技巧的前提是原数组 nums 不会发生变化。

如果原数组中的某个元素改变了，那么 preSum 数组中该元素后面的值就会失效，需要重新花费 $O(n)$ 的时间计算 preSum 数组，这就和普通的暴力解法没太大区别了。

第二个局限性：前缀和技巧只适用于存在逆运算的场景。

比方说求和的场景，你知道 $x + 6 = 10$ ，那么可以推导出 $x = 10 - 6 = 4$ ，求乘积的场景也是类似的，你知道 $x * 6 = 12$ ，那么可以推导出 $x = 12 / 6 = 2$ ，这就叫存在逆运算，都可以使用前缀和技巧。

但有些场景是没有逆运算的，比方说求最大值的场景，你知道 $max(x, 8) = 8$ ，此时你无法推导出 $x$ 的值。

想要同时解决这两个问题，就需要更高级的数据结构，最通用的解决方案是线段树，我们会在数据结构设计章节具体讲解。

303. 区域和检索 - 数组不可变 https://leetcode.cn/problems/range-sum-query-immutable

304. 二维区域和检索 - 矩阵不可变 https://leetcode.cn/problems/range-sum-query-2d-immutable

一维数组中的前缀和

二维矩阵中的前缀和

拓展延伸