树状数组

写在最前

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

树状数组也被称为二进制索引树(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);
}
};