树状数组原理

树状数组是一种可以动态维护序列前缀和的数据结构, 树状数组是用数组来模拟树形结构, 可以解决大部分基于区间上的更新以及求和问题. 功能为:

• 单点更新 update(i, v): 把序列的$i$位置的数值加上一个值$v$

• 区间查询 query(i): 查询序列$[1, 2, \cdots, i]$区间的区间和, 即位置$i$的前缀和

原理

树状数组与原数组的关系

树状数组为了节省空间, 删去了不必要的结点, 将结点数压缩到与数组长度相同. 设原数组为$A[i]$, 对应的树状数组为$C[i]$, 两个数组长度相同, 对应关系如下图所示:

$C$中的每个一位置, 都代表了一个, 或一段原始数组中的位置, 其值为所代表的原始数组片段的和. 如$C[4]$代表了$A[1]$到$A[4]$, $C[6]$代表了$A[5]$和$A[6]$, $C[8]$代表了$A[1]$到$A[8]$. 对于长度为8的数组, 完整的关系为:

• C[1] = A[1]

• C[2] = A[1] + A[2]

• C[3] = A[3]

• C[4] = A[1] + A[2] + A[3] + A[4]

• C[5] = A[5]

• C[6] = A[5] + A[6]

• C[7] = A[7]

• C[8] = A[1] + A[2] + A[3] + A[4] + A[5] + A[6] + A[7] + A[8]

当数组长度更长时, 如何确定C[i]代表那些A[i]结点呢. 我们首先要定义每个位置$i$, C[i]与数组A的对应关系.

关系探索

首先观察上图, 上图中树状数组虽然本质为数组, 但不同位置之间有着类似于树中结点之间的关系. 以树的角度来看:

• 长度为1的数组对应的树, 以C[1]为根节点, 高度为1

• 长度为2的数组对应的树, 以C[2]为根节点, 高度为2

• 长度为4的数组对应的树, 以C[4]为根节点, 高度为3

• 长度为8的数组对应的树, 以C[8]为根节点, 高度为4

• ...

观察数组长度和树高度的关系, 很容易发现存在$n=2^h$的关系, 其中$n$是数组长度, $h$是树的高度. 而相同高度的结点, 对应的原数组中元素的数量又是相同的, 例如同为高度1的C[1], C[3], C[5], C[7]都只对应原数组中一个位置, 而同为高度2的C[2]和C[6]都对应两个位置.

因此, 如果我们能找到一种映射关系, 能给定每个位置C[i]对应的高度, 也就找到了C[i]和A的对应关系.

二进制

$n=2^h$这种关系可以联想到二进制上:

• 1 -> 1

• 2 -> 10

• 4 -> 100

• 8 -> 1000

• ...

可以看到某些十进制数大小与二级制数长度满足这种关系. 进一步地, 将1~8全部转换成二进制数, 并将它们对应的高度标注在一起.

• 1 -> 1 -> 1

• 2 -> 10 -> 2

• 3 -> 11 -> 1

• 4 -> 100 -> 3

• 5 -> 101 -> 1

• 6 -> 110 -> 2

• 7 -> 111 -> 1

• 8 -> 1000 -> 4

可以看出来, 对于任意$i$, 其在树中的高度, 等于$i$对应的二进制数中, 末尾0的数量, 再加上1. 而高度又决定了关联的原数组中元素的数量:

• 1高度为1, 末尾0的数量为0, 关联的元素数量为1

• 2高度为2, 末尾0的数量为1, 关联的元素数量为2

• 3高度为1, 末尾0的数量为0, 关联的元素数量为1

• 4高度为3, 末尾0的数量为2, 关联的元素数量为4

• 5高度为1, 末尾0的数量为0, 关联的元素数量为1

• 6高度为2, 末尾0的数量为1, 关联的元素数量为2

• 7高度为1, 末尾0的数量为0, 关联的元素数量为1

• 8高度为4, 末尾0的数量为3, 关联的元素数量为8

我们把C[i]对应的A中的片段称为信息区间, 可以看到位置$i$的信息区间长度一定等于$2^{\text{i的二级制中末尾0的数量}}$, 我们将其记为$\text{lowbit}(i)=2^{\text{i的二级制中末尾0的数量}}$.

为什么信息区间的长度与二级制数字中末尾0的数量有这样关系. 末尾的0前面是一个1, 为了得到这个1, 我们需要不停的+1, 直到这个位置由0变1, 这时由于进位关系, 1后面全部为0, 而累加的1的数量就是$2^{\text{1之后0的数量}}$. 以6为例, 对应110, 我们要得到由低到高第二位的1, 需要从100开始, 添加两个1, 得到110, 所以其对应的信息区间长度为2, 得到了一棵子树. 而100的信息区间又是另一颗子树的问题.

每个高度对应的信息区间的长度有了对应关系, 以位置$i$为右端点, 以$\text{lowbit}(i)$为区间长度, 我们就得到了位置$i$的信息区间$[i - \text{lowbit}(i) + 1, i]$.

这也是为什么上图中的所有子树看起来都是朝左倾斜的, 因为$i$的信息区间要以$i$为右端点.

操作

为什么我们要设计这么一种树. 我们要得到序列每个位置的前缀和, 固然在某些位置(满足$2^n$这些数)C[i]就是对应的前缀和, 但在更多位置上是不满足的, 例如C[7]=A[7], 只包含了原数组中第7个元素的值.

而这种数据结构, 能够让我们在$O(\log{n})$的时间复杂度下, 去查询或是更新某个位置的前缀和, 非常高效. 只是需要一些特别的操作来实现.

查询

例如对于7要计算前缀和, 就是要计算A[1]到A[7]所有数字之和, C[7]只包含了A[7]的值. 但是我们知道C[7]对应长度为1的区间[7, 7]的信息, 也就找到了上一个信息区间的最右端$7 - 1 = 7 - \text{lowbit}(7) = 6$.

通过上面我们知道, C[i]就是这个信息区间元素之和. 而6对应一个长度为2的信息区间[5, 6]. 至此我们得到了区间[5, 7]的和: C[7] + C[6].

同理我们又可以找到上一个信息区间$6 - 2 = 6 - \text{lowbit}(6) = 4$. 再将C[4]加上, 得到C[7] + C[6] + C[4]. C[4]对应的信息区间为[1, 4], 再向上追溯, 有$4 - 4 = 4 - \text{lowbit}(4) = 0$, 超过了数组的范围, 停止.

因此, 借助树状数组, 查询位置$i$对应的前缀和, 就是一个不断向前寻找信息区间的工作, 直到走出数组范围, 然后将所有找到的信息区间之和加在一起, 就得到了位置$i$对应的前缀和.

而寻找上一个信息区间的迭代式为$j = i - \text{lowbit}(i)$, $j$为新的区间代表.

用程序表达为:

prefix_sum = 0
while i > 0:
    prefix_sum += C[i]
    i -= lowbit(i)

更新

更新则与查询的操作方向相反. 某个位置更新了, 要将所有包含这个位置的信息区间C[i]全部更新到, 就要向上逐级搜索父结点. 而对应的迭代式为: $j = i + \text{lowbit}(i)$, 正好与查询相反. 如此循环直到$i$大于数组长度.

其实我们要做的, 就是找到第一个使得$\text{lowbit}(j) \gt \text{lowbit}(i)$的节点$j$, 这个$j$对应的信息区间必然包含$i$代表的信息区间. $i$对应的信息区间长度为$\text{lowbit}(i)$, 我们再加上一个$\text{lowbit}(i)$, $1+1$进位得到0, 末尾0的长度增加了1, 因此得到的结点为$i + \text{lowbit}(i)$, 满足$\text{lowbit}(j) = 2 \times \text{lowbit}(i)$.

如何求解lowbit

给定$i$, 由公式快速计算得到lowbit:

$\text{lowbit}(i) = i \& (-i)$

至于为什么参考下面的文章.

• lowbit计算

用途

用来对前缀进行求和/求最大值/求最小值.

求和:

• [315][困难][线段树][二分] 计算右侧小于当前元素的个数
• [剑指Offer-51][困难][线段树][二分] 数组中的逆序对

求最大值:

• [300][中等][贪心][二分][动态规划][树状数组] 最长上升子序列

参考资料

• 树状数组的原理是什么？
• 树状数组的详细分析（含实例）