浅谈树状数组

· · 算法·理论

下面我将讲解一种非常好用的数据结构,树状数组!

前置知识

在讲解之前,我们首先来点前置知识,\text{lowbit}

什么是 \text{lowbit} 呢?我们设手里有一个数 x,那么 x\text{lowbit} 就是 x 在二进制下从右往左第一个为 1 的位代表的数。具体一点,6 在二进制下是 110 ,这个数从右往左第一个为 1 的位是从右往左第二位,这一位代表的数是 2,所以 6\text{lowbit}2。同样的,20 在二进制里表示为 10100,所以 20\text{lowbit}4

那么如何求这个 \text{lowbit} 呢?首先我们要了解 -x 的二进制。-x 在二进制下就是将 x 的二进制所有位全部取反然后 +1。这个加一是最重要的。因为只要末尾是一,加一之后就会变成零,然后进位到第二位。如果第二位也是一,继续变成零,进位。直到某一位是 0,那么就将这一位变成 1。然后就到了最关键的一步 x\And(-x)。由于取反,所以加一过程没有触及到的高位在 x-x 的二进制之中一定有一个是 0,那么与运算之后全都是 0。那么那些低位呢。在 -x 的二进制中,最后面就是我们加一运算中进位得出的 0。而进位之后是 0,原来就是 1,所以在 x 的二进制里这一位就是 0。两个都是 0,与运算之后一定是 0 了。现在就剩 -x 的二进制中从右往左第一个为 1 的数了,这一位恰好在 x 的二进制里也是 1,所以刚好可以求出 \text{lowbit(x)}

切入正题

刚刚讲了好长的前置知识,总算切入正题了。

我们首先有一个例题,就是说给定一个长度为 n 的数组 a,和 q 次询问,每次询问可以对一个数 a_l 进行加 k,或者求出 a_la_r 这段区间的和。(n,q\le 10^5)

这个数据范围,暴力求和肯定是不行的。如果只有求和,那么完全可以前缀和,如果只有修改完全可以差分。但两个都有就必须要用高级一点的数据结构,例如树状数组。

树状数组就是又有一个数组 bb_x 表示 a_{x-\text{lowbit(x)}+1}a_x 的和。然后数组 b 大概就长这样。

可以看出很像一棵树,不然为什么叫树状数组吗。

那么修改不就好说了吗。如果要修改 $a_x$,那么就需要把所有包含 $a_x$ 的 $b$ 全部修改,那就是把这棵树往上连边搜索了吗。然后我们看什么时候会有连边。容易发现当且仅当 $y=x+\text{lowbit(x)}$ 时才会有边。由于每一次操作后,下次操作就不用关心之前的二进制位了,而一个数 $n$ 只有 $\log n$ 个二进制位,所以修改的复杂度是 $O(\log n)$ 的。 下面就是我写的修改的代码。 ```cpp void cr(int oo,int sz){ while(oo<=n)b[oo]+=sz,oo+=lo(oo); } ``` 然后我们考虑查询。由于 $l$ 和 $r$ 会穿插在树状数组里,不太方便,我们考虑前缀和,先查询 $a_1$ 加到 $a_r$,在查询 $a_1$ 加的 $a_{r-1}$,相减即可。由于 $b_x$ 表示 $a_{x-\text{lowbit(x)}+1}$ 到 $a_x$ 的和,所以要查询 $a_1$ 到 $a_x$ 的和,则可以在查询完 $b_x$ 之后每次都跳到 $b_{x-\text{lowbit(x)}}$ 即可完美覆盖所有区间。同样也是每次查询到一个数之后这个二进制位已经变成 $0$ 了,下次再跳 $\text{lowbit(x)}$ 一定只会跳更左边的二进制位,每次查询复杂度也是 $O(\log n)$,已经足以通过此题。 查询的代码如下。 ```cpp int js(int oo){ int ans=0; while(oo)ans+=b[oo],oo-=lo(oo); return ans; } ``` 有了以上这些内容,就可以通过 [P3374](https://www.luogu.com.cn/problem/P3374) 了。 最后稍微留意一下在输入 $a_i$ 的时候本质上就是在做修改。 整道题的代码大概就是这样。 ```cpp #include<bits/stdc++.h> using namespace std; int n,m,a[500005],b[500005],o,x,y; int lo(int x){ return x&(-x); } void cr(int oo,int sz){ while(oo<=n)b[oo]+=sz,oo+=lo(oo); } int js(int oo){ int ans=0; while(oo)ans+=b[oo],oo-=lo(oo); return ans; } int main(){ cin>>n>>m; for(int i=1;i<=n;i++){ cin>>a[i]; cr(i,a[i]); } for(int i=1;i<=m;i++){ cin>>o; if(o==1){ cin>>x>>y; cr(x,y); a[i]+=y; } else{ cin>>x>>y; cout<<js(y)-js(x-1)<<endl; } } } ``` 然后我们稍微变一点样,变成区间修改和单点查询。 之前是区间查询,所以我们运用前缀和的思想,现在变成了区间修改,硬做肯定没法做,我们不难想到差分。 具体怎么差分呢?就是树状数组里存的已经不是数值的和了,而是差值的和。如果不理解可以看下面的图,比如 $b_6$ 就是 $(a_6-a_5)+(a_5-a_4)$。这里我们默认 $a_0$ 是 $0$。反正全局数组也是这样。 ![](https://cdn.luogu.com.cn/upload/image_hosting/4sxjg5bu.png) 好的,此时修改就太好做了。比如说让 $a_l$ 到 $a_r$ 这段区间整体加 $k$,那么我们只需要将 $a_l-a_{l-1}$ 加上 $k$,再将 $a_{r+1}-a_r$ 减去 $k$ 不就可以了吗。 至于查询呢?查询也太好做了。由于是差分数组,所以想知道第 $x$ 项,只需要知道差分数组里前 $x$ 项的和了。 这题的完整代码如下。 ```cpp #include<bits/stdc++.h> using namespace std; int n,m,a[500005],b[500005],o,x,y,z; int lo(int x){return x&(-x);} void cr(int oo,int sz){ while(oo<=n)b[oo]+=sz,oo+=lo(oo); } int js(int oo){ int ans=0; while(oo)ans+=b[oo],oo-=lo(oo); return ans; } int main(){ cin>>n>>m; for(int i=1;i<=n;i++) cin>>a[i],cr(i,a[i]-a[i-1]); for(int i=1;i<=m;i++){ cin>>o; if(o==1)cin>>x>>y>>z,cr(x,z),cr(y+1,-z); else cin>>x,cout<<js(x)<<endl; } } ``` 下一步当然就是把上两个功能结合起来了。 就是说给定一个长度为 $n$ 的数组 $a$,和 $q$ 次询问,每次询问可以对一个区间 $a_l$ 到 $a_r$ 全部加 $k$,或者求出 $a_l$ 到 $a_r$ 这段区间的和。$(n,q\le 10^5)

其实这是个线段树的板子题,如果你希望学习请看 这里,但这篇文章讲的是树状数组,所以我们用树状数组来做。

同样 a_la_r 转化为前缀和。那么差分之后单点就是前缀和了,那么前缀查询不就是前缀和的前缀和了吗。假设查询 a_1a_x 的和那么我们就可以考虑,第一个数字被 x 个前缀和覆盖,第二个数被 x-1 个前缀和覆盖,如此类推。最后弄出来的式子就是这样。

\sum_{i=1}^{x}a_i\times(x-i+1)

展开后得

\sum_{i=1}^{x}a_i\times x+\sum_{i=1}^{x}a_i-\sum_{i=1}^{x}a_i\times i

那么就维护两个树状数组,第一个存 a_i,第二个存 a_i\times i,就可以完美解决,我写的代码就是这样。

#include<bits/stdc++.h>
#define int long long
using namespace std;
int n,m,a[500005],b[500005],c[500005],o,x,y,z;
int lo(int x){return x&(-x);}
void cr(int oo,int sz){
    while(oo<=n)b[oo]+=sz,oo+=lo(oo);
}
int js(int oo){
    int ans=0;
    while(oo)ans+=b[oo],oo-=lo(oo);
    return ans;
}
void cr2(int oo,int sz){
    while(oo<=n)c[oo]+=sz,oo+=lo(oo);
}
int js2(int oo){
    int ans=0;
    while(oo)ans+=c[oo],oo-=lo(oo);
    return ans;
}
signed main(){
    cin>>n>>m;
    for(int i=1;i<=n;i++)
    cin>>a[i],cr(i,a[i]-a[i-1]),cr2(i,(a[i]-a[i-1])*i);
    for(int i=1;i<=m;i++){
        cin>>o;
        if(o==1)cin>>x>>y>>z,cr(x,z),cr(y+1,-z),cr2(x,z*x),cr2(y+1,-z*(y+1));
        else cin>>x>>y,cout<<(js(y)*y+js(y)-js2(y))-(js(x-1)*(x-1)+js(x-1)-js2(x-1))<<endl;
    }
}

所以树状数组常数小,耗时短,耗空间小,代码短,能有树状数组的题为何偏要写线段树呢?

上述所有问题都是树状数组里存的是数值,但通过下面问题你可以发现,树状数组里可以存的不是数值。

给定一个数 n 和一个数组 a,要求 a 中逆序对的数量。

大家在学归并排序思想时肯定做过,但是这题也可以用树状数组做,而且思维难度更低。我们不能盲目枚举所有的逆序对,但是我们可以用树状数组存储每一个数出现的次数。什么意思,比如说,这里出现了一个 6,那我们就把树状数组原本第六个位置加一,之后从这里往上跑 \text{lowbit(x)},每跑到一个位置都加一。然后修改完之后立刻查询。查询的是对于这个数 a_i,以它为结尾的逆序对数量。

什么时候才会出现逆序对?显然是前面的数比后面的数大。那么我们就看前面有几个数比这个数大不就可以了吗。虽然树状数组板子是绿但是明明这么做更简单。然后放一下我的代码。

#include<bits/stdc++.h>
#define int long long
using namespace std;
int n,m,o,x,y,ans;
map<int,int>a,b;
int lo(int x){return x&(-x);}
void cr(int oo,int sz){
    while(oo<=1e9+1)b[oo]+=sz,oo+=lo(oo);
}
int js(int oo){
    int ans=0;
    while(oo)ans+=b[oo],oo-=lo(oo);
    return ans;
}
signed main(){
    cin>>n;
    for(int i=1;i<=n;i++){
        cin>>a[i],cr(a[i],1),ans+=i-js(a[i]);
    }
    cout<<ans;
}

顺便警示后人一下,修改函数的上限千万不要写成 n 了,要不然你会对着代码调很久还不知道哪里错了。

例题选讲

刚刚讲的东西都是模版,从这里进入例题选讲阶段。

AT_abc441_e

题目传送门

这题同样有别的解法,但是我们这里讲树状数组。

如果直接算,感觉没啥头绪。我们可以考虑转化成前缀 AB 多的数量。我们可以发现,如果 lr 这一子串 AB 多,那么 1l-1AB 多的数量一定比 1r 的要少。

那么就变成了一个正序对的板子。

由于代码和刚刚的逆序对实在没啥大区别所以不放了。唯一需要注意的就是树状数组里不能存负数,为了防止 AB 少,我们将树状数组里所有存的数都加上一个固定的数,这样不影响大小。

AT_abc436_f

题目传送门

我们看这题,发现直接枚举位置不太好做。那么考虑枚举别的东西,比如说枚举最暗的星星,然后任意区间里比这颗星星暗的一定不能出现,比这颗星星亮的必须出现。

那就好说了,对于每一颗星星,它左边比它亮的星星数量加一再乘上它右边它亮的星星数量加一就是答案。

对于为什么要加一,因为区间可以正好卡到它那里,左边或右边就一颗星星也拍不到了。

然后把此题我的代码也放上来。

#include<bits/stdc++.h>
#define int long long
using namespace std;
int n,m,o,x,y,ans,a[500005],b[500005],c[500005];
int lo(int x){return x&(-x);}
void cr(int oo,int sz){while(oo<=n)b[oo]+=sz,oo+=lo(oo);}
int js(int oo){
    int ans=0;
    while(oo)ans+=b[oo],oo-=lo(oo);
    return ans;
}
void cr2(int oo,int sz){while(oo<=n)c[oo]+=sz,oo+=lo(oo);}
int js2(int oo){
    int ans=0;
    while(oo)ans+=c[oo],oo-=lo(oo);
    return ans;
}
signed main(){
    cin>>n;
    for(int i=1;i<=n;i++)
        cin>>a[i],cr(a[i],1);
    for(int i=1;i<=n;i++){
        ans+=(js2(a[i]-1)+1)*(js(a[i]-1)-js2(a[i]-1)+1),cr2(a[i],1);
    }
    cout<<ans;
}

P10814

这是一个二维数点题,一看题目就很像树状数组,但你可能会被一个问题卡住。就是说树状数组只能说从左往右边存储边计算的,这里的 lr 没有保证单调递增啊。

但是,没关系,因为这题没有强制要求在线。这题的做法就是先按 l-1 排序,从左到右处理完所有的之后再按 r 排序,继续处理。都处理完之后,我们就排序排回来,然后前缀和作差即可。至于怎么处理,上面已经讲的很详细了,你可以尝试自己去写了。

顺便讲一下什么是离线和在线。在线就是说每输入一组数据就能立即输出。离线就是说把所有数据输出完之后一起统计。有的时候题目会让你用这组数据输入的数异或上上一组数据的答案再进行计算,这样就强制在线了。

同样放一下我的代码呀。

#include<bits/stdc++.h>
using namespace std;
int n,m,a[2000005],b[2000005],c[2000005],ans[2000005],cnt;
struct node{
    int x,y,z,id,x2,y2,fl;
}d[4000005];
bool cmp(node xx,node yy){return xx.x<yy.x;}
bool cnp(node xx,node yy){return xx.y<yy.y;}
bool czp(node xx,node yy){return xx.id<yy.id;}
int lo(int x){return x&(-x);}
void cr(int oo,int sz){while(oo<=2e6)b[oo]+=sz,oo+=lo(oo);}
void cr2(int oo,int sz){while(oo<=2e6)c[oo]+=sz,oo+=lo(oo);}
int js(int oo){
    int ans=0;
    while(oo)ans+=b[oo],oo-=lo(oo);
    return ans;
}
int js2(int oo){
    int ans=0;
    while(oo)ans+=c[oo],oo-=lo(oo);
    return ans;
}
int main(){
    int x,y,z;
    cin>>n>>m;
    for(int i=1;i<=n;i++)
    cin>>a[i];
    for(int i=1;i<=m;i++){
        cin>>x>>y>>z;
        d[++cnt].x=x-1;
        d[cnt].z=z;
        d[cnt].id=i;
        d[cnt].fl=0;
        d[++cnt].x=y;
        d[cnt].z=z;
        d[cnt].id=i;
        d[cnt].fl=1;
    }
    sort(d+1,d+cnt+1,cmp);
    for(int i=0,jl=1;i<=n;i++){
        if(i!=0)cr(a[i],1);
        while(d[jl].x==i){
            ans[d[jl].id]+=(d[jl].fl?1:-1)*js(d[jl].z),jl++;
        }
    }
    for(int i=1;i<=m;i++)
    cout<<ans[i]<<'\n';
}

总结

树状数组是一种非常好用的数据结构,代码段,常数低。无论是做单点修改区间查询,区间修改单点查询,还是区间修改区间查询,都很实用。在做逆序对,二维数点等题时同样适用。非常建议好好学学。本文中也提到了几个写树状数组代码时需注意的细节,同样非常重要的。

希望你们看完之后都能收获满满!别忘了点个赞!