树状数组

写在最前

树状数组算是我目前学过最优雅好写的数据结构了。

树状数组也被称为二进制索引树(Binary Indexed Tree,BIT)，是一种高效的数据结构，它可以实现 $O(\log n)$ 的单点修改和前缀和查询。

核心思想

当我们想实现单点修改和区间查询时，肯定会有人想到用一个数组来维护若干个小区间的和，使得修改和查询的复杂度都不会那么高。

树状数组便是巧妙的运用了二进制来实现这一想法。

现在我们存储每一个位置前面 lowbit 位的和，也就是是说 $c_i=\sum_{j=i-lowbit(i)+1}^ia[j]$，比如 $i=6$，那么我们可以先将 6 转化为二进制 $(0110)_2$，显然 $c[6]=a[5]+a[6]$。

实现

区间查询

这里偷借一下老师的图（（（（

不难发现长得确实很像树如果现在要求前 $i$ 项的和，那么就可以通过每次减去 $lowbit(i)$ 得到的 $c_i$ 求和，具体来说就是 $sum[i]=c[i]+sum[i-lowbit[i]]$。

还是拿 6 举例，先化成二进制 $(0110)_2$。

  • $lowbit((0110)_2)=(0010)_2$，即 $lowbit(6)=2$
  • $(0110)_2-(0010)_2=(0100)_2$，即 $6-2=4$
  • $lowbit((0100)_2)=(0100)_2$，即 $lowbit(4)=4$
  • $(0100)_2-(0100)_2=(0000)_2$，即 $4-4=0$

所以 $c[6]=c[4]+c[2]$。

代码：

1
2
3
4
5
6
7
8
9
10
long long query(long long x)
{
    long long sum=0;
    while(x)
    {
        sum+=c[x];
        x-=(x&-x);
    }
    return sum;
}

单点修改

查询是从上往下退，修改也同理，从下往上爬就行了，不难发现，$i$ 的父节点即为 $i+lowbit(i)$。

代码：

1
2
3
4
5
6
7
8
void update(long long x,long long v)
{
    while(x<=n)
    {
        c[x]+=v;
        x+=(x&-x);
    }
}

不过这样实在是太屎山了，不妨把他封装起来：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
class BIT
{
    long long c[N];
    long long prefix(long long x)
    {
        long long sum=0;
        while(x)
        {
            sum+=c[x];
            x-=(x&-x);
        }
        return sum;
    }
public:
    void update(long long x,long long v)
    {
        while(x<=n)
        {
            c[x]+=v;
            x+=(x&-x);
        }
    }
    long long query(long long l,long long r)
    {
        return prefix(r)-prefix(l-1);
    }
}
...
BIT c;//定义
c.update(i,x)//第 I 位加 x
c.query(l,r)//查询 l~r 的和

进阶

区间修改单点查询

结合单点修改区间查询，不难想到差分，我们只需要对差分数组做树状数组，这样需要区间修改的话只需要单点修改差分数组的两端，单点修改的话需要修改一整个区间。

1
2
c.update(l,x);c.update(r+1,-x);//l~r 加 x
c.query(1,x)//查询第 x 位

区间修改区间查询

核心思想

仍然是差分，通过差分数组 $d$，我们可以将区间修改转化为单点修改：

  • 将区间 $[l, r]$ 内的所有元素加上 $x$，等价于：
    • $d[l] += x$
    • $d[r+1] -= x$

同时，区间查询可以通过差分数组的前缀和来实现：

  • $
    a[i] = \sum_{j=1}^i d[j]
    $
  • $\sum_{i=1}^k a[i] = \sum_{i=1}^k \sum_{j=1}^i d[j] = \sum_{j=1}^k d[j] \cdot (k - j + 1) = (k + 1) \cdot \sum_{j=1}^k d[j] - \sum_{j=1}^k d[j] \cdot j$

为了高效计算上述公式，我们需要维护两个树状数组：

  1. 一个树状数组维护差分数组 $d[i]$。
  2. 另一个树状数组维护 $d[i] \cdot i$。

实现细节

修改

对于区间 $[l, r]$ 的修改操作，我们需要更新两个树状数组：

  • 在第一个树状数组中：
    • $d[l] += x$
    • $d[r+1] -= x$
  • 在第二个树状数组中：
    • $d[l] \cdot l += x \cdot l$
    • $d[r+1] \cdot (r+1) -= x \cdot (r+1)$

代码实现：

1
2
3
4
5
6
void update(int l, int r, long long x) {
    modify(d, l, x);
    modify(d, r + 1, -x);
    modify(id, l, x * l);
    modify(id, r + 1, -x * (r + 1));
}
查询

对于区间 $[1, k]$ 的查询操作，我们可以通过以下公式计算：
$\sum_{i=1}^k a[i] = (k + 1) \cdot \text{query}d(k) - \text{query}{id}(k)$其中：

  • $\text{query}d(k)$ 是第一个树状数组的前缀和，即 $\sum{i=1}^k d[i]$。
  • $\text{query}{id}(k)$ 是第二个树状数组的前缀和，即 $\sum{i=1}^k d[i] \cdot i$。

代码实现：

1
2
3
4
5
6
long long query(int l, int r) {
    auto get_sum = [&](int k) {
        return (k + 1) * query(d, k) - query(id, k);
    };
    return get_sum(r) - get_sum(l - 1);
}

喜闻乐见封装环节

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
class BIT {
private:
    int n;
    long long d[N];
    long long id[N];
    void modify(int x, long long v, long long arr[]) {
        while (x <= n) {
            arr[x] += v;
            x += (x & -x);
        }
    }
    long long ask(int x, long long arr[]) {
        long long sum = 0;
        while (x) {
            sum += arr[x];
            x -= (x & -x);
        }
        return sum;
    }
    long long get_sum(int k) {
        return (k + 1) * ask(k, d) - ask(k, id);
    }

public:
    BIT(int size) {
        n = size;
        for (int i = 0; i <= n + 1; i++) {
            d[i] = 0;
            id[i] = 0;
        }
    }
    void update(int l, int r, long long x) {
        modify(l, x, d);
        modify(r + 1, -x, d);
        modify(l, x * l, id);
        modify(r + 1, -x * (r + 1), id);
    }
    long long query(int l, int r) {
        return get_sum(r) - get_sum(l - 1);
    }
};