树状数组笔记

11jiang08 / 2023-08-21 / 原文

树状数组是算法竞赛中常见的高级数据结构，是利用数的二进制特征进行检索的树状数据结构。常用于求解动态变化的区间问题。具有结构清晰，操作方便，应用灵活，效率高的特点。

单点修改+区间查询

这是树状数组最基础的应用，设 \(BIT[i]\) 表示 \([i-lowbit(i)+1,i]\) 区间内的数字和。此处 sum 的作用是求区间 \([1,x]\) 值的和。

void LSB(int x){return x&(-x);}
void add(int x,int d){
    while(x<=n){BIT[x]+=d;x+=LSB(x);}
}
int query(int x){
    int ans=0;
    while(x){ans+=BIT[x];x-=LSB(x);}
    return ans;
}

区间修改+单点查询

可以利用差分数组将区间修改变为单点修改。设 \(d[i]\) 为 \(a[i]\) 的差分数组，修改区间 \([L,R]\) 时，将 \(d[L]\) 增加，\(d[R+1]\) 减少，求解 \(a[x]\) 即求 \(\displaystyle \sum _{i=1}^x d[i]\)。可以利用树状数组维护。

void LSB(int x){return x&(-x);}
void add(int x,int d){
    while(x<=n){BIT[x]+=d;x+=LSB(x);}
}
void change(int l,int r,int d){add(l,d);add(r+1,-d);}
int query(int x){
    int ans=0;
    while(x){ans+=BIT[x];x-=LSB(x);}
    return ans;
}

区间修改+区间查询

这里需要两个树状数组才能高效完成 “区间修改+区间查询”，称为 “二阶树状数组”。

定义 \(d[i]\) 为 \(a[i]\) 的差分数组，区间修改的方法上面已经讲过，下面推导区间和：

\(\displaystyle \sum _{i=1}^k a_k=d_1+(d_1+d_2)+...+ \displaystyle \sum _{i=1}^k d_k\)

\(=k·d_1+(k-1)·d_2+...+d_k\)

\(=k\displaystyle \sum_{i=1}^k d_i-\sum_{i=1}^k(i-1)d_i\)

因此我们可以维护两个树状数组，一个计算 \(d_i\) 的前缀和，一个计算 \((i-1)·d_i\) 的前缀和。在 \(O(\log n)\) 的时间内进行 “区间修改+区间查询”。

void add(int x,int d1,int d2){
    while(x<=n){
        BIT1[x]+=d1;  BIT2[x]+=d2;
        x+=LSB(x);
    }
}
int query1(int x){
    int ans=0;
    while(x){ans+=BIT1[x];x-=LSB(x);}
    return ans;
}
int query2(int x){
    int ans=0;
    while(x){ans+=BIT2[x];x-=LSB(x);}
    return ans;
}
void update(int l,int r,int d){
    add(l,d,d*(l-1));
    add2(r+1,-d,-d*r);
}
void query(int l,int r){
    return r*query1(r)-query2(r)-(l-1)*query(l-1)+query2(l-1);
}

二维 “区间修改+区间查询”

二维 “区间修改+区间查询” 是一维的拓展，推导方法和一维的类似。

令 \(d[i][j]\) 为 \(a[i][j]\) 的差分，有：

\[d[i][j]=a[i][j]-a[i-1][j]-a[i][j-1]+a[i-1][j-1] \]

此时 \(a[i][j]=\displaystyle \sum_{k=1}^i \sum _{l=1}^jd[k][l]\)。

接下来推导区间和，以顶点为 \((a,b)\) 和 \((c,d)\) 的矩阵的区间和，有：

\[\small\displaystyle \sum_{i=a}^c\sum_{j=b}^d a[i][j]=\sum_{i=1}^c\sum_{j=1}^d a[i][j]-\sum_{i=1}^c\sum_{j=1}^{b-1} a[i][j]+\sum_{i=1}^{a-1}\sum_{j=1}^d a[i][j]+\sum_{i=1}^{a-1}\sum_{j=1}^{b-1} a[i][j] \]

可以看出问题的关键在于求解 \(\displaystyle\sum_{i=1}^n\sum_{j=1}^m a[i][j]\) 的值。

\(\displaystyle\sum_{i=1}^n\sum_{j=1}^m a[i][j]\)

\(=\displaystyle\sum_{i=1}^n\sum_{j=1}^m \sum_{k=1}^i \sum _{l=1}^jd[k][l]\)

\(=\displaystyle\sum_{i=1}^n\sum_{j=1}^m d[i][j]·(n-i+1)·(m-j+1)\)

\(\small=(n+1)(m+1)\displaystyle\sum_{i=1}^n\sum_{j=1}^m d[i][j]-(m+1)\sum_{i=1}^n\sum_{j=1}^m d[i][j]·i-(m+1)\sum_{i=1}^n\sum_{j=1}^m d[i][j] · j+\sum_{i=1}^n\sum_{j=1}^m d[i][j]·ij\)

此时需要维护四个树状数组，分别计算 \(d[i][j],d[i][j]·i,d[i][j]·j,d[i][j]·ij\)

void add(int x,int y,int d){
    for(int i=x;i<=n;i+=LSB(i)){
        for(int j=y;j<=m;j+=LSB(j)){
            t1[i][j]+=d;  t2[i][j]+=x*d;
            t3[i][j]+=y*d;  t4[i][j]+=x*y*d;
        }
    }
}
void update(int a,int b,int c,int d,int delta){
    add(a,b,delta);  add(c+1,d+1,delta);
    add(a,d+1,-delta);  add(c+1,b,-delta);
}
int sum(int x,int y){
    int ans=0;
    for(int i=x;i;i-=LSB(i)){
        for(int j=y;j;j-=LSB(j)){
            ans+=(x+1)*(y+1)*t1[i][j]-(y+1)*t2[i][j]-(x+1)*t3[i][j]+t4[i][j];
        }
    }
    return ans;
}
void query(int a,int b,int c,int d){
    return sum(c,d)+sum(a-1,b-1)-sum(c,b-1)-sum(a-1,d);
}