浅谈树状数组

树状数组是一种功能类似线段树的数据结构,但在代码量上较线段树更少,也同样能进行对序列的修改与查询操作。

什么是树状数组?

树状数组, 或 Binary Indexed Tree(也称 Fenwick Tree ),是一种能够高效更新并计算数列前缀和的数据结构。

树状数组的实现主要基于计算二进制 lowbit(也就是值为 11 的最低位)。

原理

既然树状数组与线段树类似,我们不妨将线段树的实现方式先搬过来看看。对于一个表示数列 aia_i 的线段树 SS ,线段树会通过数列 aa 递归建树,使用类似归并的思想维护节点。一开始,线段树为空,然后从根节点开始分枝,再进行递归展开,如此往复,直到叶子结点为止。

也就是说,在线段树中每一个非叶子结点都管理着当前孩子结点的信息。这种操作的具体实现就是 pushup 操作。这样在查询时,我们只需要找到对应的父亲节点即可。

而树状数组所使用的策略与线段树几乎别无二致:同样适用父亲节点来表示其孩子节点的信息,且同为递归定义。同样地,在查询时只需要查询所属叶子结点的父亲节点即可。显然,对于根节点将会管理数列中的所有节点的信息。

fenwick-1

上图就是一个典型的使用树状数组维护出来的树形结构。最上方的 a[] 表示数据数组,一一对于最下方的叶子结点(圆圈)。每一个父亲节点都使用方形表示,控制着至少两个结点。

将上图中的结点进行编号,得:

fenwick-2

每个节点右面的蓝色数字为当前节点的编号。如果将树状数组表达为 bib_i ,原始数组为 aia_i ,则有:

b1=a1b2=a1+a2b3=a1+a2+a3b4=a4b5=a4+a5b6=a4+a5+a6 b_1 = a_1\\ b_2 = a_1 + a_2\\ b_3 = a_1 + a_2 + a_3\\ b_4 = a_4\\ b_5 = a_4 + a_5\\ b_6 = a_4 + a_5 + a_6
如果把下标都转换成二进制,则有:
b(1)2=a(1)2b(10)2=a(1)2+a(10)2b(11)2=a(1)2+a(10)2+a(11)2b(100)2=a(100)2b(101)2=a(100)2+a(101)2b(110)2=a(100)2+a(101)2+a(110)2 b_{(1)_2} = a_{(1)_2}\\ b_{(10)_2} = a_{(1)_2} + a_{(10)_2}\\ b_{(11)_2} = a_{(1)_2} + a_{(10)_2} + a_{(11)_2}\\ b_{(100)_2} = a_{(100)_2}\\ b_{(101)_2} = a_{(100)_2} + a_{(101)_2}\\ b_{(110)_2} = a_{(100)_2} + a_{(101)_2} + a_{(110)_2}
不难发现其中的规律:
bi=ai2k+1+ai2k+2++ai b_i = a_{i - 2^k + 1} + a_{i - 2^k + 2} + \cdots + a_i
其中 kkii 的二进制表示下 11 的最低位的位置。例如 i=(8)10=(1000)2i = (8)_{10} = (1000)_2 时,k=3, b8=a1+a2++a8k = 3,\ b_8 = a_1 + a_2 + \cdots +a_8

考虑计算下述表达式:

si=j=1iaj s_i= \sum_{j=1}^{i} a_j
根据 (3) 式易得:
si=bi+bi2k1+bi2k12k2+ s_i = b_i + b_{i - 2^{k_1}} + b_{i - 2^{k_1} - 2^{k_2}} + \cdots
其中 kjk_j 表示当前第 jj 次访问时,从 ii 一直循环对最低位上的 11 取反后当前 ii 的值的 11 的最低位。例如,对于计算 s4s_4 ,我们有:
s4=b4+b420=b4+b3=a4+a1+a2+a3=25 s_4 = b_4 + b_{4 - 2^0} = b_4 + b_3 = a_4 + a_1 + a_2 + a_3 = 25
其中 kk 的二进制变化过程为:
k=(4)10=(100)2k=(100)2(100)2=(0)2End. k = (4)_{10} = (100)_2 \rarr k = (100)_2 - (100)_2 = (0)_2 \rarr \text{End.}
这样我们就完成了前缀求和操作。

对于修改操作,我们也可以采用类似的思想。修改操作无非就是将给定的增加值依次向上更新传递。例如,将 a5a_5 增加 33 我们需要依次向上更新 b5b_5b6b_6 。这样的变化过程在求和操作时已经说明,只需要进行反向操作即可。

对于 a5=a(101)2a_5 = a_{(101)_2} ,我们需要更新其所有祖宗节点,即所有包括 a5a_5 值的 bib_i 。我们不妨从 i=(5)10=(101)2i = (5)_{10} = (101)_2 进行操作。注意到 a5a_5 的父亲节点包括 {b5,b6}={b(101)2,b(110)2}\{b_5, b_6\} = \{b_{(101)_2}, b_{(110)_2}\} 。所以可以将 ii 每次增加 2k2^k ,其中 kkii 的二进制表示下 11 的最低位的位置。i=5i = 5kk 的变化过程如下:

i=5=(101)2, k=0i=i+2k=5+1=6, k=1End. i = 5 = (101)_2,\ k = 0 \rarr i = i + 2^k = 5 + 1 = 6,\ k = 1 \rarr \text{End.}
这样我们就可以进行单点修改操作。

实现

注意到上述所有操作都依赖于一个重要操作:求某数 xx 的二进制下最低 11 的位置 kk。这时不妨引出一个简单的函数 lowbit

int lowbit(int x) {
  return x & -x;
}

该函数的返回值即为 2k2^k ,而树状数组所需操作也仅仅是 2k2^k 的值。该函数原理如下:

  1. x=(17)10x = (17)_{10} ,则其补码为 (010001)(010001)_补(x)(-x)_补 = (101111)(101111)_补
  2. 计算 x & -x ,结果为 (000001)2=(1)10(000001)_2 = (1)_{10}
  3. 验证上述操作,xx 的二进制最低 11 的位置从零开始计数为 002k=20=12^k = 2^0 = 1 ,结果正确。

于是我们便不难写出模板代码:

template <typename T = int>
class FenwickTree {
   private:
    T tr[N];

    int lowbit(int x) { return x & -x; }

   public:
    void add(int x, T k) {
        while (x <= N) {
            tr[x] += k;
            x += lowbit(x);
        }
    }

    T prefix_sum(int x) {
        T ans = 0;
        while (x >= 1) {
            ans += tr[x];
            x -= lowbit(x);
        }
        return ans;
    }
  
    T sum(int l, int r) {
      return prefix_sum(r) - prefix_sum(l - 1);
    }
};

不难发现虽然基本操作是前缀和,但仍然可以使用两个前缀和的差求出任意区间和。值得注意的是,虽然 lowbit 函数仅支持整数操作,但树状数组支持任意数据类型:lowbit 维护的是数组下标,并不对实际内容进行操作。

不难发现求和与修改操作都是 O(logn)O(\log n) 的:每次操作的最坏情况即为从叶子结点更新至根节点,也就是需要遍历 h=树高h = \text{树高} 次。而树状数组又是完全二叉树,所以 hlog2(n)h \leq \log_2(n)

对比线段树,我们发现树状数组并不能很高效地处理区间更新问题。注意到原本的前缀和操作,我们不妨将其改为差分操作。这样,每次存储 bib_i 则为 aa 的差分数组,即 bi=di+di1++d1b_i = d_i + d_{i - 1} + \cdots + d_1 ,其中 did_i 表示 aa 的差分数组,即 di=aiai1d_i = a_i - a_{i - 1}

这样,对于修改操作我们只需要对 blb_l 增加 xxbrb_r 减少 xx 即可。当前这样的坏处就是无法进行区间查询操作:对差分数组进行求和相对原数组而言已经失去意义。但是,我们仍可以通过差分数组的特性来查询某一节点的值:对前 ii 个差分数组元素进行求和即为原数组的第 ii 个元素的值。

尽管无法直接进行区间求和的维护操作,我们仍可以通过维护两个树状数组的方式同时进行区间修改与区间查询的操作。参考 OI-Wiki 。但是本人认为进行这样的操作不如使用线段树方便。

例题:https://www.luogu.com.cn/problem/P3368

显然对于这道题,我们需要进行区间修改与单点查询。参考代码如下:

// Problem: P3368 【模板】树状数组 2
// Contest: Luogu
// URL: https://www.luogu.com.cn/problem/P3368
// Memory Limit: 125 MB
// Time Limit: 1000 ms
//
// Powered by CP Editor (https://cpeditor.org)

#include <memory.h>

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <iostream>
#include <queue>
#include <stack>
#include <string>
#include <vector>
using namespace std;
const int N = 5 * 1e5 + 10;

int bit[N], n, m, x, y, k, op, a[N];

int lowbit(int x) { return x & -x; }

void add(int x, int k) {
    while (x <= n) {
        bit[x] += k;
        x += lowbit(x);
    }
}

int prefix_sum(int x) {
    int ans = 0;
    while (x >= 1) {
        ans += bit[x];
        x -= lowbit(x);
    }
    return ans;
}

int main() {
    cin >> n >> m;
    for (int i = 1; i <= n; i++) {
        cin >> a[i];
    }
    for (int i = 1; i <= n; i++) {
        add(i, a[i] - a[i - 1]);
    }
    for (int i = 1; i <= m; i++) {
        cin >> op;
        if (op == 1) {
            cin >> x >> y >> k;
            add(x, k);
            add(y + 1, -k);
        } else {
            cin >> x;
            cout << prefix_sum(x) << endl;
        }
    }
    return 0;
}

注意到核心模板代码无需改变。

总结

不难发现,朴素树状数组可以进行大部分线段树可以进行的操作:区间求和与单点更新、区间修改与单点查询。但是,这两组操作却无法使用一个树状数组同时实现。简单来说,树状数组能做的操作线段树也一定能做,但反之却不亦然。

在对于高级操作没有要求的题目下,使用树状数组是一个很好的选择,因为其代码复杂度明显较低。但在复杂的题目时,线段树会略胜一筹。

值得注意的是,虽然线段树与树状数组时间复杂度相同,但树状数组常数略小。在更新操作时,树状数组不一定一定要遍历 logn\log n 个节点,而线段树却必须遍历所有 logn\log n 个节点。但是该常数复杂度可以忽略,二者时间消耗均等。但是树状数组在空间消耗上要明显低于线段树:线段树通常需要 44 倍的数组空间,而树状数组仅需要 11 倍。

总之,哪个好用用哪个。