树状数组

正文索引 [隐藏]

本期我们来讲一讲树状数组,一个支持单点修改、区间求和的数据结构。(经过精巧的修改后,也可以支持区间修改&查询、单点修改&区间最大值)

假设我们有一个长度为 n 的序列 a_1, a_2, \ldots, a_n,我们有两个操作:

  1. 修改操作:选择序列中的一个数字 a_k,给它增加为 x
  2. 查询操作:选择序列的一个区间 [l, r],求区间内元素之和 \sum_{i=l}^r a_i

朴素做法

对于这个问题,让我们先了解两个较为朴素的做法。

暴力算法

既然有这样一个序列,那么我们就申请一个数组来维护序列中每一个元素的值。
来了一个修改操作,我们直接修改对应数组的值。
来了一个查询操作,我们遍历区间内的所有元素求一个和。

这样做,修改操作的复杂度是 O(1),查询操作的复杂度是 O(n)。如果修改操作非常多,而查询操作几乎没有,那么算法就运行的非常快,反之很慢。

前缀和算法

既然查询慢,那么我们就优化查询操作。
我们现在不使用数组维护序列的元素,而是维护序列的元素和。
具体来说我们让数组的第 k 个元素 arr[k] = \sum_{i=1}^k a_i
现在来了一个查询操作 [l, r],我们可以直接通过 arr[r] – arr[l – 1] 得到结果,这样查询的复杂度就降低到了 O(1)
但是这种方法的问题在于修改操作,如果我们要修改一个元素 a_i,那么我们就要修改求和时计算了 a_i 的所有数组元素 arr[i], arr[i+1], \ldots, arr[n]。于是修改操作的时间复杂度上升成了 O(n)

这种方法叫做前缀和,可以发现前缀和的修改非常慢,而查询非常快,正好和暴力算法相反。

树状数组

暴力和前缀和,在这个问题上正好是两个极端,暴力的修改快、查询慢,前缀和的查询快、修改慢。
暴力的修改快是因为一个数组元素就对应一个序列元素的值,修改只牵扯单个数组元素,而查询需要遍历区间内所有的数组元素,故很慢。
前缀和则相反,因为序列中一个元素会被很多数组元素所统计,所以查询的时候很容易就可以凑出序列和,而修改的时候则很痛苦。

我们是否能从这两个算法中获得启发,得到一个修改较快、查询也较快的算法呢?
想法就是我们要找到序列到数组的统计关系,使得凑出任意序列和不需要太多的数组元素,同时任意一个序列元素不会被太多数组所统计。

具体方法

首先必须说明的是树状数组没有直接查询任意区间和的能力,它能做到的仅仅是查询任意前缀和,而区间和需要通过两个前缀和相减得到。
树状数组做到了:对于任意一个前缀和,树状数组只需要取出 O(\log n) 界别的元素就可以凑出它,当我们要修改序列某一个元素时,也仅仅牵扯到 O(\log N) 个数组元素。

树状数组的内在原理还是二进制,这里我们先举一个例子:
当给我们一个前缀和查询 \sum_{i=1}^k a_i 的时候,我们通过二进制将这个查询分成 \log_2 k 个数组元素组合得到。
假设 k = 21_{(10)} = 10101_{(2)},我们根据 k 二进制上的 1 找到 如下几个数组元素:

\begin{aligned}
arr[10101_{(2)}] &= \sum_{i=10100_{(2) + 1}}^{10101_{(2)}} a_i\\
arr[10100_{(2)}] &= \sum_{i=10000_{(2) + 1}}^{10100_{(2)}} a_i\\
arr[10000_{(2)}] &= \sum_{i=00000_{(2) + 1}}^{10000_{(2)}} a_i\\
\end{aligned}

那么,很容易就通过这三个数组维护的元素和得到了 \sum_{i=1}^{10101_{(2)}} a_i 的答案。

这里我们定义一个函数 lowbit(x) 可以求得 x 数位上最低的二进制位上的 1,例如,lowbit(10100_{(2)}) = 100_{(2)}

数组第 k 个位置,其实维护了序列中 (k-lowbit(k), k] 位置的元素,数组到序列元素的维护关系是这样的:

arr[k] = sum_{i = k – lowbit(k) + 1}^{k} a_i

在有了这样的维护关系之后,对于任意一个前缀和查询操作,我们都可以像上面的例子一样,通过二进制位快速找到如何凑出答案。