LeetCode 例题精讲 | 18 前缀和：空间换时间的技巧
LeetCode 例题精讲 | 18 前缀和：空间换时间的技巧
本文将教会你「前缀和」的算法套路，做出以下 LeetCode 例题：
- LeetCode 724. Find Pivot Index（Easy）
- LeetCode 560. Subarray Sum Equals K 和为K的子数组（Medium）
在设计算法时，时间复杂度始终是我们关注的重点。我们需要让算法的时间复杂度尽可能低，追求运行效率。有些时候，我们可以通过增加空间占用的方式减少算法的运行时间，这便是空间换时间。
动态规划就是一类空间换时间的算法。动态规划通过保存所有子问题的计算结果，可以避免子问题的重复计算。这种方法的代价是 DP 数组 占用了较多的空间。
前缀和同样也是一种空间换时间的技巧，只不过我们使用的不是 DP 数组，而是「前缀和数组」。
那么，究竟什么是前缀和呢？
什么是前缀和
我们先用一道简单的题目理解一下「前缀和」究竟是做什么的：
LeetCode 303. Range Sum Query - Immutable（Easy）
这道题目的解法很直白，难点在于如何减少时间复杂度。我们来看看不同的解法的时间、空间复杂度有何区别。
解法一：暴力法
解法一：暴力法
public int sumRange(int i, int j) {
int sum = 0;
for (int k = i; k <= j; k++) {
sum += nums[k];
}
return sum;
}
解法二：空间换时间
在 sumRange 会被调用很多次的情况下，我们要尽可能地减少一次调用的时间。如果多次调用 sumRange 的参数是重复的，但还需要重新求和，就会做很多重复的计算。
为了避免重复的计算，我们可以对数组 nums 进行预处理，预先存储计算结果。我们使用二维数组 res 存储预处理的结果，res[i][j] 存储 sumRange(i, j) 的返回值。
解法二：空间换时间
private int[][] res;
// 预处理阶段
public NumArray(int[] nums) {
int n = nums.length;
res = new int[n][n];
for (int i = 0; i < n; i++) {
int sum = 0;
for (int j = i; j < n; j++) {
sum += nums[j];
res[i][j] = sum;
}
}
}
public int sumRange(int i, int j) {
return res[i][j];
}
解法三：前缀和
上一个解法中的预处理方法过于暴力，会空间占用太大。我们还可以使用另一个更聪明的预处理方法：前缀和。
所谓前缀和（prefix sum），就是数组开头的若干连续元素的和。
前缀和的定义
语言小贴士： 前缀和数组的定义中使用了左闭右开区间。这种表示方法的优点之一是很容易做区间的减法。例如：
nums[0..j) - nums[0..i)可以得到nums[i..j)。在滑动窗口类题目中也经常使用左闭右开区间。
sumRange(i, j) = preSum[j+1] - preSum[i]
int n = nums.length;
// 计算前缀和数组
int[] preSum = new int[n+1];
preSum[0] = 0;
for (int i = 0; i < n; i++) {
preSum[i+1] = preSum[i] + nums[i];
}
最终的题解代码如下所示：
class NumArray {
private int[] preSum;
// 预处理阶段
public NumArray(int[] nums) {
int n = nums.length;
// 计算前缀和数组
preSum = new int[n+1];
preSum[0] = 0;
for (int i = 0; i < n; i++) {
preSum[i+1] = preSum[i] + nums[i];
}
}
public int sumRange(int i, int j) {
return preSum[j+1] - preSum[i];
}
}
我们可以对比一下三种解法的时间、空间复杂度。
可以看到，前缀和方法的特点是：能优化时间复杂度，同时让空间复杂度不会太大。这让前缀和成为一个很实用的数组预处理手段。
前缀和的应用
下面，我们用两道典型题目来看看前缀和的应用场景。这两道题分别是「寻找枢纽元素」以及「和为K的子数组」。
前缀和方法的典型使用场景是数组类题目。当看到题目与「子数组求和」有关，就要想想能不能使用前缀和来求解。
例题一：寻找枢纽元素
LeetCode 724. Find Pivot Index 寻找枢纽元素（Easy）
示例：
输入：nums = [1, 7, 3, 6, 5, 6]
输出：3
解释：索引 3 (nums[3] = 6) 的左侧数之和 (1 + 7 + 3 = 11)，与右侧数之和 (5 + 6 = 11) 相等。
对于题目中定义的「枢纽元素」，我们可以用一张图来理解：
枢纽元素的含义
public int pivotIndex(int[] nums) {
// 首先计算所有元素之和 S
int S = 0;
for (int n : nums) {
S += n;
}
int A = 0; // A 为前缀和
// 迭代计算前缀和
for (int i = 0; i < nums.length; i++) {
int x = nums[i];
if (2 * A + x == S) {
// 满足公式中的关系，x 是枢纽元素
return i;
}
A += x; // 计算前缀和
}
return -1;
}
例题二：和为K的子数组
LeetCode 560. Subarray Sum Equals K 和为K的子数组（Medium）
这道题关注就是「子数组的元素之和」，显然又是一道可以使用前缀和技巧的题目。
我们可以首先求出所有的前缀和，然后根据前缀和求出所有可能的子数组之和，题解代码如下所示：
public int subarraySum(int[] nums, int k) {
int N = nums.length;
// 计算前缀和数组
// presum[k] 表示元素 nums[0..k) 之和
int[] presum = new int[N+1];
int sum = 0;
for (int i = 0; i < N; i++) {
presum[i] = sum;
sum += nums[i];
}
presum[N] = sum;
// sum of nums[i..j) = sum of nums[0..j) - sum of nums[0..i)
int count = 0;
for (int i = 0; i <= N; i++) {
for (int j = i+1; j <= N; j++) {
// 前缀和相减求子数组之和
if (presum[j] - presum[i] == k) {
count++;
}
}
}
return count;
}
注意： 哈希表技巧不是本文的重点，这里只是介绍本题的最优解法。关于哈希表的相关技巧在后面的文章中会有专门的讲解。
我们再仔细看一下上面代码中的二重循环：
for (int i = 0; i <= N; i++) {
for (int j = i+1; j <= N; j++) {
// 前缀和相减求子数组之和
if (presum[j] - presum[i] == k) {
count++;
}
}
}
为了减少时间复杂度，我们的目标是把二重循环变为一重循环。
我们将条件判断 presum[j] - presum[i] == k 简单移项，可以得到 presum[i] == presum[j] - k。那么我们的循环完全可以把 i、j 颠倒，写成这样：
for (int j = 1; j <= N; j++) {
for (int i = 0; i < j; i++) {
if (presum[i] == presum[j] - k) {
count++;
}
}
}
public int subarraySum(int[] nums, int k) {
// 前缀和 -> 该前缀和（的值）出现的次数
Map<Integer, Integer> presum = new HashMap<>();
// base case，前缀和 0 出现 1 次
presum.put(0, 1);
int sum = 0; // 前缀和
int res = 0;
for (int n : nums) {
sum += n; // 计算前缀和
// 查找有多少个 sum[i] 等于 sum[j] - k
if (presum.containsKey(sum - k)) {
res += presum.get(sum - k);
}
// 更新 sum[j] 的个数
if (presum.containsKey(sum)) {
presum.put(sum, presum.get(sum) + 1);
} else {
presum.put(sum, 1);
}
}
return res;
}
总结
本文介绍了前缀和的技巧，以及相关的两道例题：LeetCode 724. 寻找枢纽元素、LeetCode 560. 和为K的子数组。这两道题目都有一个共同点：对子数组求和。我们在做题的时候，只要遇到与「子数组求和」相关的题目，就考虑一下使用前缀和方法会如何，一定没有错。
前缀和技巧掌握起来并不难，但是需要一定的经验才能在题目中灵活运用。还没有用过前缀和的同学，建议自己做一遍这两道例题，体会一下前缀和在时间复杂度的优化。
前缀和是一个「小方法」，在解题过程中往往是进行局部的优化，一般还需要和其他方法一起使用。例如 560 题就还需要利用哈希表做进一步的优化，以消除不必要的循环。与哈希表相关的技巧将在后续的文章中进一步介绍。