前缀和数组: 快速计算数组区间和

Motivations

有这么一类问题——给定一个数组 $arr$ 和 $Q$ 个查询，每一个查询的格式是 $query(l, r)$，意思是计算区间和 $arr[l] + arr[l + 1] + … + arr[r]$

如果采用暴力求解，那么查询区间和的时间复杂度是 $O(N)$，处理 $Q$ 个查询就需要 $O(NQ)$，有没有什么数据结构或者是算法可以优化这个时间复杂度呢？

前缀和数组就可以做到，这就是今天要讲的主题

前缀和数组定义

先看前缀和数组，定义如下

$$ \texttt{prefix}[k]=\sum_{i=0}^{k-1}arr[i] $$

换言之，$\texttt{prefix}[k]$ 是数组 $arr$ 前 $k$ 个数字的和，这也是为什么它叫做前缀和数组

计算出前缀和数组 $\texttt{prefix}$ 只需要 $O(N)$ 的时间复杂度，按顺序遍历 $arr$ 数组，利用如下的公式

$$ \texttt{prefix}[k+1] = \texttt{prefix}[k] + arr[k] $$

根据前缀和数组的定义，区间和的计算可以利用如下的式子

$$ query[l, r]=\texttt{prefix}[r+1]-\texttt{prefix}[l] $$

可以看到，现在处理一个区间查询的时间复杂度变成了 $O(1)$，处理 $Q$ 个查询的时间复杂度是

$$ O(N + Q)\ll O(NQ) $$

前缀和数组实现

按照定义写如下

prefix_sum = [0 for _ in range(len(arr) + 1)]
for i in range(len(arr)):
	prefix_sum[i + 1] = prefix_sum[i] + arr[i]

当然我更经常使用的是如下的写法 :)

from itertools import accumulate
prefix_sum = list(accumulate(a, initial=0))

Lazy Update

前面我们的数组是 $arr$ 是静态的，即读取进来之后就不会再更改了。但有时候做题会发现，题目要求我们支持这么一种更新操作 $update(arr, l, r, x)$——对区间 $[l, r]$ 的数字批量加上一个数字或者减去一个数字。一道典型的题目是是 SPOJ.HAYBALE-Haybale stacking

还是先来分析一下暴力算法的时间复杂度。假设有 $K$ 个这样的更新操作，处理单个 $update(arr, a, b, x)$ 操作的时间复杂度是 $O(N)$，那么 $K$ 个更新操作就是 $O(NK)$，显然时间复杂度太高了，在算法题里面经常都会 TLE

为了节省性能开销，不难想到不应该真的去更新，而是应该惰性更新（Lazy Update），那么要怎么做呢？答案就在下面这个公式上

$$ \texttt{prefix}[i+1] = \texttt{prefix}[i] + arr[i] $$

前缀和数组会不断累加原数组的值，这里存在一个“累加”的效果，试想一下，如果我们让 $arr[l] + x$ 会怎么样？可想而知，从 $l$ 这个位置开始的前缀和数组都会 $+x$，到这里答案已经呼之欲出了，我们再让 $arr[r + 1] - x$，从 $r + 1$ 这个位置开始的前缀和数组都会 $-x$。等价的效果就是 $[l, r]$ 这个区间里面的元素都 $+x$ 而不影响其他位置的值。Python 代码大致如下

def update(arr, l, r, x):
    arr[l] += x
    arr[r + 1] -= x

这样更新的操作就是 $O(1)$ 了， $K$ 个更新操作只需要 $O(K)$，最后再用 $O(N)$ 计算前缀和数组，总的时间复杂度为 $O(N + K)$

$$ O(N + K)\ll O(NK) $$