前缀和技巧适用于快速、频繁地计算一个索引区间内的元素之和。
参考labuladong算法小抄:小而美的算法技巧:前缀和数组
一维数组中的前缀和
力扣第 303 题「区域和检索 - 数组不可变」,让你计算数组区间内元素的和,这是一道标准的前缀和问题:. - 力扣(LeetCode)
【python 前缀和解法】
class NumArray:
def __init__(self, nums: List[int]):
# pre_sum数组首位为0,故pre_sum长度比nums多1
self.pre_sum = [0]
for i in range(len(nums)):
# pre_sum[i] = pre_sum[i-1] + nums[i-1]
self.pre_sum.append(self.pre_sum[i] + nums[i])
def sumRange(self, left: int, right: int) -> int:
return self.pre_sum[right+1] - self.pre_sum[left]
以上代码构造了一个首位为0,长度比nums多1的pre_sum数组,旨在后续求sumRange方便计算和减少判断逻辑。若构造一个和nums等长度的pre_sum也可以,只是在求sumRange时遇到区间左侧left为0时需要多一个判断逻辑。后续二维矩阵汇中pre_matrix比matrix多一行一列也是类似道理。
class NumArray:
def __init__(self, nums: List[int]):
# pre_sum首位即为nums[0],长度与nums相等
self.pre_sum = [nums[0]]
for i in range(1, len(nums)):
self.pre_sum.append(self.pre_sum[i-1] + nums[i])
def sumRange(self, left: int, right: int) -> int:
# 注意left为0的特殊情况
return self.pre_sum[right] - (self.pre_sum[left-1] if left != 0 else 0)
【思路】
题目要求你实现这样一个类:
class NumArray {
public NumArray(int[] nums) {}
/* 查询闭区间 [left, right] 的累加和 */
public int sumRange(int left, int right) {}
}
sumRange
函数需要计算并返回一个索引区间之内的元素和,没学过前缀和的人可能写出如下代码:
class NumArray {
private int[] nums;
public NumArray(int[] nums) {
this.nums = nums;
}
public int sumRange(int left, int right) {
int res = 0;
for (int i = left; i <= right; i++) {
res += nums[i];
}
return res;
}
}
这样,可以达到效果,但是效率很差,因为sumRange
方法会被频繁调用,而它的时间复杂度是O(N)
,其中N
代表nums
数组的长度。
这道题的最优解法是使用前缀和技巧,将sumRange
函数的时间复杂度降为O(1)
,说白了就是不要在sumRange
里面用 for 循环,咋整?
直接看代码实现:
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];
}
}
## 核心思路是我们 new 一个新的数组preSum
出来,preSum[i]
记录nums[0..i-1]
的累加和,看图 10 = 3 + 5 + 2:
看这个preSum
数组,如果我想求索引区间[1, 4]
内的所有元素之和,就可以通过preSum[5] - preSum[1]
得出。
这样,sumRange
函数仅仅需要做一次减法运算,避免了每次进行 for 循环调用,最坏时间复杂度为常数O(1)
。
二维矩阵中的前缀和
力扣第 304 题「304. 二维区域和检索 - 矩阵不可变」,其实和上一题类似,上一题是让你计算子数组的元素之和,这道题让你计算二维矩阵中子矩阵的元素之和:
【python 二维矩阵-前缀和解法】
class NumMatrix:
def __init__(self, matrix: List[List[int]]):
m, n = len(matrix), (len(matrix[0]) if matrix else 0)
# 创建长为m+1,宽为n+1的pre_matrix,为后续算法做铺垫
self.pre_matrix = [[0] * (n+1) for _ in range(m+1)]
# 浅拷贝
pre_matrix = self.pre_matrix
# 关键注意:pre_matrix的长和宽分别为m+1和n+1
for i in range(1,m+1):
for j in range(1,n+1):
pre_matrix[i][j] = pre_matrix[i][j-1] + pre_matrix[i-1][j] -\
pre_matrix[i-1][j-1] + matrix[i-1][j-1]
def sumRegion(self, x1: int, y1: int, x2: int, y2: int) -> int:
pre_matrix = self.pre_matrix
return pre_matrix[x2+1][y2+1] - pre_matrix[x2+1][y1] -\
pre_matrix[x1][y2+1] + pre_matrix[x1][y1]
【思路】
比如说输入的matrix
如下图:
那么sumRegion([2,1,4,3])
就是图中红色的子矩阵,你需要返回该子矩阵的元素和 8。
这题的思路和一维数组中的前缀和是非常类似的,如下图:
如果我想计算红色的这个子矩阵的元素之和,可以用绿色矩阵减去蓝色矩阵减去橙色矩阵最后加上粉色矩阵,而绿蓝橙粉这四个矩阵有一个共同的特点,就是左上角就是(0, 0)
原点。
那么我们可以维护一个二维preSum
数组,专门记录以原点为顶点的矩阵的元素之和,就可以用几次加减运算算出任何一个子矩阵的元素和:
class NumMatrix {
// preSum[i][j] 记录矩阵 [0, 0, i, j] 的元素和
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];
}
}
这样,sumRegion
函数的复杂度也用前缀和技巧优化到了 O(1)。
和为 k 的子数组
最后聊一道稍微有些困难的前缀和题目,力扣第 560 题「和为 K 的子数组」:
【python 前缀和进阶解法】
class Solution:
def subarraySum(self, nums, k):
preSum = 0
freq = {}
cnt = 0
for v in nums:
# 先存入前一个位置的前缀和个数,恰好也预先计入了 freq[0] = 1
freq[preSum] = freq.get(preSum, 0) + 1
# 后更新当前位置前缀和,该前缀和此时未计入 freq,限制 r' < r
preSum += v
cnt += freq.get(preSum - k, 0)
return cnt
【思路】
那我把所有子数组都穷举出来,算它们的和,看看谁的和等于k
不就行了,借助前缀和技巧很容易写出一个解法:
int subarraySum(int[] nums, int k) {
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];
int res = 0;
// 穷举所有子数组
for (int i = 1; i <= n; i++)
for (int j = 0; j < i; j++)
// 子数组 nums[j..i-1] 的元素和
if (preSum[i] - preSum[j] == k)
res++;
return res;
}
这个解法的时间复杂度O(N^2)
空间复杂度O(N)
,并不是最优的解法。不过通过这个解法理解了前缀和数组的工作原理之后,可以使用一些巧妙的办法把时间复杂度进一步降低。
注意前面的解法有嵌套的 for 循环:
for (int i = 1; i <= n; i++)
for (int j = 0; j < i; j++)
if (preSum[i] - preSum[j] == k)
res++;
第二层 for 循环在干嘛呢?翻译一下就是,在计算,有几个j
能够使得preSum[i]
和preSum[j]
的差为k
。毎找到一个这样的j
,就把结果加一。
我们可以把 if 语句里的条件判断移项,这样写:
if (preSum[j] == preSum[i] - k)
res++;
优化的思路是:我直接记录下有几个preSum[j]
和preSum[i] - k
相等,直接更新结果,就避免了内层的 for 循环。我们可以用哈希表,在记录前缀和的同时记录该前缀和出现的次数。
int subarraySum(int[] nums, int k) {
int n = nums.length;
// map:前缀和 -> 该前缀和出现的次数
HashMap<Integer, Integer>
preSum = new HashMap<>();
// base case
preSum.put(0, 1);
int res = 0, sum0_i = 0;
for (int i = 0; i < n; i++) {
sum0_i += nums[i];
// 这是我们想找的前缀和 nums[0..j]
int sum0_j = sum0_i - k;
// 如果前面有这个前缀和,则直接更新答案
if (preSum.containsKey(sum0_j))
res += preSum.get(sum0_j);
// 把前缀和 nums[0..i] 加入并记录出现次数
preSum.put(sum0_i,
preSum.getOrDefault(sum0_i, 0) + 1);
}
return res;
}
比如说下面这个情况,需要前缀和 8 就能找到和为k
的子数组了,之前的暴力解法需要遍历数组去数有几个 8,而优化解法借助哈希表可以直接得知有几个前缀和为 8。
这样,就把时间复杂度降到了O(N)
,是最优解法了。
前缀和技巧就讲到这里,应该说这个算法技巧是会者不难难者不会,实际运用中还是要多培养自己的思维灵活性,做到一眼看出题目是一个前缀和问题。