前缀和(Prefix Sum)算法笔记C++
前缀和(Prefix Sum)算法介绍
前缀和是一种预处理技术, 用于快速计算数组中任意子区间的元素之和. 它通过一次遍历创建一个辅助数组来存储从数组起始位置到当前索引位置所有元素的累加和, 从而使得后续查询操作的时间复杂度降低至 O ( 1 ) O(1) O(1).
定义
对于给定的数组 nums
, 它的前缀和数组 prefixSum
是一个新的数组, 其中 prefixSum[i]
表示原数组 nums
中从 nums[0]
到 nums[i]
的所有元素的和. 特别地, prefixSum[0] = nums[0]
.
计算过程
- 初始化
prefixSum[0] = arr[0]
. - 对于每个
i > 0
, 设置prefixSum[i] = prefixSum[i-1] + arr[i]
.
计算前缀和数组通常通过一次遍历原数组来完成. 以 C++ 代码为例, 实现如下:
std::vector<int> num{
1, 2, 3, 4, 5};
std::vector<int> prefix_sum(num.size());
prefix_sum[0] = num[0];
for (int i = 1; i < num.size(); ++i) {
prefix_sum[i] = prefix_sum[i - 1] + num[i];
}
fmt::println("num : {}", num);
fmt::println("prefix sum: {}", prefix_sum);
输出:
num : [1, 2, 3, 4, 5]
prefix sum: [1, 3, 6, 10, 15]
或者使用 C++标准库的std::partial_sum
函数:
std::vector<int> num{
1, 2, 3, 4, 5};
std::vector<int> prefix_sum(num.size());
std::partial_sum(num.begin(), num.end(), prefix_sum.begin());
fmt::println("num : {}", num);
fmt::println("prefix sum: {}", prefix_sum);
代码输出同上.
应用场景
快速计算子数组的和: 一旦得到了前缀和数组, 就可以在 O(1)
的时间复杂度内计算出原数组中任意子数组 [left, right]
的和.
子数组 [left, right]
的和为 prefixSum[right] - (prefixSum[left - 1] if left > 0 else 0)
.
例如, 对于上述 num
数组, 如果要计算子数组 [2, 4]
的和, 即 prefixSum[4] - prefixSum[1] = (1 + 2 + 3 + 4 + 5) - (1 + 2) = 15 - 3 = 12
.
优化某些算法: 在一些算法中, 通过使用前缀和可以将原本时间复杂度较高的算法优化为更高效的算法. 例如, 在计算数组中所有子数组的和时, 使用前缀和可以避免重复计算, 从而降低时间复杂度.
例题
LeetCode 560. 和为 K 的子数组
给定一个整数数组 nums
和一个整数 k
, 返回数组中和为 k
的子数组的个数.
样例:
-
nums = [1, 1, 1], k = 2
: 结果为 2, 因为[1, 1]
与[1, 1]
是两组和为2
的子数组. -
nums = [1, 2, 3], k = 3
: 结果为 2, 因为[1, 2]
与[3]
是两组和为3
的子数组. -
nums = [1, -1, 1, 2, 3, -2, -1], k = 3
: 结果为 7, 情况如图所示:
从例三中, 我们可以看到对于一个 Index
i
, 以它结尾的子数组可能不止一个.