【模板】树状数组

1.概述

树状数组（Binary Indexed Trees），其基本用途是维护序列的前缀和。对于给定的序列 $a[]$，我们建立一个数组 $c[]$，其中 $c[x]$ 保存序列 $a[]$ 的区间 $[x-lowbit(x)+1,x]$ 中所有数的和。

事实上，数组 $c[]$ 可以看作一个如下图所示的树形结构，图中最下边一行是 $N$ 个叶节点 $(N=16)$，代表数值 $a[1\sim N]$。该结构满足以下性质：

• 每个内部节点 $c[x]$ 保存以它为根的子树中所有叶节点的和。

• 每个内部节点 $c[x]$ 的子节点个数等于 $lowbit(x)$ 的位数。

• 除树根外，每个内部节点 $c[x]$ 的父节点是 $c[x+lowbit(x)]$。

• 树的深度为 $O(\log N)$。

• 如果 $N$ 不是 $2$ 的整次幂，那么树状数组就是一个具有同样性质的森林结构。

2.原理

(1) 对于一个普通的二叉树，叶子结点代表数组 $a[]$ 的 $a[1]\sim a[8]$。

(2) 将其进行简单的变形。

(3) 然后定义每一列的顶端结点数组 $c[]$，令 $c[i]$ 代表子树的叶结点的权值之和。

• $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]$

(4) 将数组 $c[]$ 的结点序号转化为二进制。

• $1=(001),c[1]=a[1]$

• $2=(010),c[2]=a[1]+a[2]$

• $3=(011),c[3]=a[3]$

• $4=(100),c[4]=a[1]+a[2]+a[3]+a[4]$

• $5=(101),c[5]=a[5]$

• $6=(110),c[6]=a[5]+a[6]$

• $7=(111),c[7]=a[7]$

• $8=(1000),c[8]=a[1]+a[2]+a[3]+a[4]+a[5]+a[6]+a[7]+a[8]$

可以发现 $c[i]=a[i-2^k+1]+a[i-2^k+2]+......+a[i]$，其中，$k$ 为 $i$ 的二进制中从最低位到高位连续零的长度。

可见，树状数组根本上来说就是二进制的应用。

要注意的是，树状数组只能计算 $a[1]$ 开始的和，$a[0]$ 这个元素是不能用的。

3.实现

3.1 lowbit(x)

$lowbit(x)$ 是 $x$ 的二进制表达式中最低位的 $1$ 所对应的值，换言之，$lowbit(x)=2^k$，$k$ 为 $i$ 的二进制中从最低位到高位连续零的长度。

对于一个数 $x$，想求 $lowbit(x)$，可借助 $x$ 的负数 $-x$，进行按位与运算。

• $x=6=(0110)$，此时 $k=1$

• $-x=-6=(1010)$

• $x$ & $(-x)=(0010)=2^1$

// 返回x的二进制最右边1所对应的值
int lowbit(int x)
{
    return x & -x;
}

$lowbit(x)$ 在树状数组中的应用：

$$\begin{array}{rl}c[i]& =a[i-2^k+1]+a[i-2^k+2]+......+a[i]\\ & =a[i-lowbit(i)+1]+a[i-lowbit(i)+2]+......+a[i]\end{array}$$

3.2 查询前缀和

树状数组支持的基本操作有两个，第一个操作是查询前缀和，即序列 $a[]$ 中第 $1\sim x$ 个数的和。

首先求出 $x$ 的二进制表示中每个等于 $1$ 的位，然后把 $[1,x]$ 分成 $O(\log N)$ 个小区间，而每个小区间的区间和都已经保存在数组 $c[]$ 中了。

下面代码在 $O(\log N)$ 的时间内查询前缀和：

// 返回a[1]+...+a[x]的和
int ask(int x)
{
    int res = 0;
    for (int i = x; i; i -= lowbit(i))
    {
        res += c[i];
    }
    return res;
}

举例：求 $a[1]+a[2]+...+a[7]$ 的和

• $7(111)$，res += c[7]

• $7-lowbit(7)=7-lowbit(111)=7-1=6(110)$，res += c[6]

• $6-lowbit(6)=6-lowbit(110)=6-2=4(100)$，res += c[4]

• $4-lowbit(4)=4-lowbit(100)=4-4=0(000)$

举例：求 $a[1]+a[2]+...+a[5]$ 的和

• $5(101)$，res += c[5]

• $5-lowbit(5)=5-lowbit(101)=5-1=4(100)$，res += c[4]

• $4-lowbit(4)=4-lowbit(100)=4-4=0(000)$

当然，若要查询序列 $a[]$ 的区间 $[l,r]$ 中所有数的和，只需计算 ask(r) - ask(l - 1)。

3.3 单点增加

树状数组支持的第二个基本操作是单点增加，意思是给序列中的一个数 $a[x]$ 加上 $y$，同时正确维护序列的前缀和。

根据上面给出的树形结构和它的性质，只有节点 $c[x]$ 及其所有祖先节点保存的“区间和”包含 $a[x]$，而任意一个节点的祖先至多只有 $\log N$ 个，我们逐一对它们的 $c$ 值进行更新即可。

下面代码在 $O(\log N)$ 的时间内执行单点增加操作：

// 令a[x] += y
void add(int x, int y)
{
    for (int i = x; i <= n; i += lowbit(i))
    {
        c[i] += y;
    }
}

如图所示，当更新 $a[1]$ 时，需要向上更新 $c[1]$、$c[2]$、$c[4]$、$c[8]$，则有：

• $1(001)$，c[1] += a[1]

• $1+lowbit(1)=1+lowbit(001)=1+1=2(010)$，c[2] += a[1]

• $2+lowbit(2)=2+lowbit(010)=2+2=4(100)$，c[4] += a[1]

• $4+lowbit(4)=4+lowbit(100)=4+4=8(1000)$，c[8] += a[1]

可以发现：更新过程是查询过程的逆过程。

4.初始化

在执行所有操作之前，我们需要对树状数组进行初始化，即针对原始序列 $a[]$ 构造一个树状数组 $c[]$。

为了简便，比较一般的初始化方法是：直接建立一个全为 $0$ 的数组 $c[]$，然后对每个位置 $x$ 执行 add(x, a[x])，就完成了对原始序列 $a[]$ 构造树状数组的过程，时间复杂度为 $O(N\log N)$。通常采用这种初始化方法已经足够。

更高效的初始化方法是：从小到大依次考虑每个节点 $x$，借助 $lowbit$ 运算扫描它的子节点并求和。若采用这种方法，上面树形结构中的每条边只会被遍历一次，时间复杂度为 $O(N)$。