前缀和(Prefix Sum)算法笔记C++

前缀和(Prefix Sum)算法介绍

前缀和是一种预处理技术, 用于快速计算数组中任意子区间的元素之和. 它通过一次遍历创建一个辅助数组来存储从数组起始位置到当前索引位置所有元素的累加和, 从而使得后续查询操作的时间复杂度降低至 $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 的子数组的个数.

样例:

1. nums = [1, 1, 1], k = 2: 结果为 2, 因为 [1, 1] 与 [1, 1] 是两组和为 2 的子数组.

2. nums = [1, 2, 3], k = 3: 结果为 2, 因为 [1, 2] 与 [3] 是两组和为 3 的子数组.

3. nums = [1, -1, 1, 2, 3, -2, -1], k = 3: 结果为 7, 情况如图所示:
   

   从例三中, 我们可以看到对于一个 Index i, 以它结尾的子数组可能不止一个.

思路<