数据结构-树状数组专题（2）

一、前言

接上回树状数组专题（1），这次主要介绍差分跟树状数组联动实现区间更新

二、我的模板

重新放了一遍，还是提一嘴，注意下标从0开始，区间左闭右开

cpp 复制代码

template <typename T>
struct Fenwick {
    int n;
    std::vector<T> a;

    Fenwick(int n_ = 0) {
        init(n_);
    }

    void init(int n_) {
        n = n_;
        a.assign(n, T{});
    }

    void add(int x, const T &v) {
        for (int i = x + 1; i <= n; i += i & -i) {
            a[i - 1] = a[i - 1] + v;
        }
    }

    T sum(int x) {
        T ans{};
        for (int i = x; i > 0; i -= i & -i) {
            ans = ans + a[i - 1];
        }
        return ans;
    }

    T rangeSum(int l, int r) {
        return sum(r) - sum(l);
    }

    int select(const T &k) {
        int x = 0;
        T cur{};
        for (int i = 1 << std::__lg(n); i; i /= 2) {
            if (x + i <= n && cur + a[x + i - 1] <= k) {
                x += i;
                cur = cur + a[x - 1];
            }
        }
        return x;
    }
};

三、差分简介

差分实际上是前缀和的逆运算，类似于函数求微积分实际上是函数求导函数的逆运算。

因此我们可以得到一个结论

差分数组的前缀和等于原数组，原数组的前缀和就等于前缀和数组

前缀和的差分等于原数组，原数组的差分就等于差分数组

我们回忆一下前缀和是怎么来的

prefix $i$ = prefix $i-1$ +a $i$

那么移项可以得到

a $i$ = prefix $i$ -prefix $i-1$

这就是差分数组的构造方式，此时的a数组就相当于差分数组，prefix数组相当于原数组

因此就有

cpp 复制代码

for(int i = 1;i <= n;++i){
    prefix[i] = prefix[i-1]+a[i];
    diff[i] = a[i]-a[i-1];
}

其中diff数组就是a数组的差分数组

区间修改

因此为了实现区间修改，如果要把 $l,r$ 这个区间的数都加上x，那么我们可以对构造出来的差分数组做这样的两个操作：

diff $l$ +=x;diff $r+1$ -=x;

这时候你要单点查询左端点l的值，实际上就是对diff $1$ 到diff $l$ 求前缀和，由于我们树状数组的特性，可以在O(logn)的时间复杂度下完成求前缀和，你会发现a $l$ 的的确确加上了x，再依次对 $l,r$ 中的值实验，发现确实都增加了x，从r+1开始，后面的值还是保持原状不增不减，这就达到了我们需要的效果

原理分析：

diff $1$ = a $1$ - a $0$

diff $2$ = a $2$ - a $1$

diff $3$ = a $3$ - a $2$

...

diff $l$ = a $l$ - a $l-1$

...

diff $r$ = a $r$ - a $r-1$

diff $r+1$ = a $r+1$ - a $r$

...

diff $n$ = a $n$ - a $n-1$

对diff $l$ 做前缀和实际上等价于求diff $1$ +diff $2$ +diff $3$ +...+diff $l$

=(a $1$ -a $0$ )+(a $2$ -a $1$ )+(a $3$ -a $2$ )+...+(a $l$ -a $l-1$ )=a $l$ -a $0$ =a $l$

因此如果diff $l$ +=x的话

整体求前缀和就等价于diff $1$ +diff $2$ +diff $3$ +...+(diff $l$ +x)

=(a $1$ -a $0$ )+(a $2$ -a $1$ )+(a $3$ -a $2$ )+...+(a $l$ -a $l-1$ +x)=a $l$ +x-a $0$ =a $l$ +x

同理，也可验证diff $r+1$ +x的情况

这里就不过多阐述了

因此我们用两个O(1)的操作实现了区间修改，取代了暴力修改的O(n)的时间复杂度

区间查询

由于我们采用差分数组来构造的树状数组，那么我们要求原数组修改后的区间和，原本是需要对差分数组做两遍前缀和的，但是我们这里可以尝试构造出两个树状数组来实现求区间和

如果要求 $l,r$ 的区间和，我们就需要先学会计算前缀和，假设我们要求前k个数的前缀和

那么对于差分数组diff\[\]来说，我们枚举出来，应该需要计算

a $1$ = diff $1$

a $2$ = diff $1$ + diff $2$

a $3$ = diff $1$ + diff $2$ + diff $3$

...

a $k$ = diff $1$ + diff $2$ + diff $3$ + ... + diff $k$

似乎找不到什么有用的结论，我们不妨添加一行，再将剩余位置补全

a $0$ = diff $1$ + diff $2$ + diff $3$ + ... + diff $k$

a $1$ = diff $1$ + diff $2$ + diff $3$ + ... + diff $k$

a $2$ = diff $1$ + diff $2$ + diff $3$ + ... + diff $k$

a $3$ = diff $1$ + diff $2$ + diff $3$ + ... + diff $k$

...

a $k$ = diff $1$ + diff $2$ + diff $3$ + ... + diff $k$

这时候我们发现，我们要求的a $1$ 到a $k$ 的和，就等于这个矩形的和减去红色部分的和

因此前k个数的和（前缀和）为(k+1)*fenwick1.sum(k)-fenwick2.sum(k)

fenwick1表示第一个树状数组，fenwick2表示第二个树状数组

fenwick1中保存的是数组a的差分数组diff，fenwick2中保存的是i*diff $i$ (1<=i<=n)

这块还是以会写代码为主

四、专题训练

4.1 星码StarryCoding P41 树状数组（区间修改）

思路

这题是树状数组区间修改的模板题，主要要会写代码

我的代码

cpp 复制代码

#include <bits/stdc++.h>
using i64 = long long;
constexpr int N = 2e5+5;

template <typename T>
struct Fenwick {
    int n;
    std::vector<T> a;

    Fenwick(int n_ = 0) {
        init(n_);
    }

    void init(int n_) {
        n = n_;
        a.assign(n, T{});
    }

    void add(int x, const T &v) {
        for (int i = x + 1; i <= n; i += i & -i) {
            a[i - 1] = a[i - 1] + v;
        }
    }

    T sum(int x) {
        T ans{};
        for (int i = x; i > 0; i -= i & -i) {
            ans = ans + a[i - 1];
        }
        return ans;
    }

    T rangeSum(int l, int r) {
        return sum(r) - sum(l);
    }

    int select(const T &k) {
        int x = 0;
        T cur{};
        for (int i = 1 << std::__lg(n); i; i /= 2) {
            if (x + i <= n && cur + a[x + i - 1] <= k) {
                x += i;
                cur = cur + a[x - 1];
            }
        }
        return x;
    }
};

int main() {
    std::cin.tie(nullptr)->sync_with_stdio(false);

    int n,q;
    std::cin >> n >> q;
    std::vector<i64> a(N,0),diff1(N,0),diff2(N,0);
    for(int i = 1;i <= n;++i) {
        std::cin >> a[i];
        diff1[i] = a[i]-a[i-1];
        diff2[i] = i*diff1[i];
    }
    Fenwick<i64> fenwick1(N),fenwick2(N);
    for(int i = 1;i <= n;++i) {
        fenwick1.add(i-1,diff1[i]);
        fenwick2.add(i-1,diff2[i]);
    }
    while(q--) {
        int op,l,r;i64 x;
        std::cin >> op >> l >> r;
        if(op==1) {
            std::cin >> x;
            fenwick1.add(l-1,x);
            fenwick2.add(l-1,l*x);
            fenwick1.add(r,-x);
            fenwick2.add(r,(r+1)*(-x));
        }else {
            i64 pre1 = fenwick1.sum(r)*(r+1)-fenwick2.sum(r);
            i64 pre2 = fenwick1.sum(l-1)*l-fenwick2.sum(l-1);
            std::cout << pre1-pre2 << '\n';
        }
    }
    
    return 0;
}

4.2 AcWing 242. 一个简单的整数问题

思路

跟上一题如出一辙，这个题还更加简单，因为只需要实现区间修改和单点查询，这就意味着我们只需要用一个树状数组就可以解决

我的代码

cpp 复制代码

#include <bits/stdc++.h>
using i64 = long long;
constexpr int N = 1e5+5;

template <typename T>
struct Fenwick {
    int n;
    std::vector<T> a;

    Fenwick(int n_ = 0) {
        init(n_);
    }

    void init(int n_) {
        n = n_;
        a.assign(n, T{});
    }

    void add(int x, const T &v) {
        for (int i = x + 1; i <= n; i += i & -i) {
            a[i - 1] = a[i - 1] + v;
        }
    }

    T sum(int x) {
        T ans{};
        for (int i = x; i > 0; i -= i & -i) {
            ans = ans + a[i - 1];
        }
        return ans;
    }

    T rangeSum(int l, int r) {
        return sum(r) - sum(l);
    }

    int select(const T &k) {
        int x = 0;
        T cur{};
        for (int i = 1 << std::__lg(n); i; i /= 2) {
            if (x + i <= n && cur + a[x + i - 1] <= k) {
                x += i;
                cur = cur + a[x - 1];
            }
        }
        return x;
    }
};

int main(){
    std::cin.tie(nullptr)->sync_with_stdio(false);
    
    int n,m;
    std::cin >> n >> m;
    std::vector<int> a(N,0);
    for(int i = 1;i <= n;++i){
        std::cin >> a[i];
    }
    std::vector<int> diff(N,0);
    for(int i = 1;i <= n;++i){
        diff[i] = a[i] - a[i-1];
    }
    Fenwick<int> fenwick(N);
    for(int i = 1;i <= n;++i){
        fenwick.add(i-1,diff[i]);
    }
    while(m--){
        std::string op;
        int l,r,d,x;
        std::cin >> op;
        if(op=="Q"){
            std::cin >> x;
            int ans = fenwick.sum(x);
            std::cout << ans << '\n';
        }else{
            std::cin >> l >> r >> d;
            fenwick.add(l-1,d);
            fenwick.add(r,-d);
        }
    }
    
    return 0;
}

4.3 AcWing 243. 一个简单的整数问题2

思路

跟4.1换个题目场景罢了，解题思路一模一样，甚至代码改几个地方就能拿过来AC掉

我的代码

cpp 复制代码

#include <bits/stdc++.h>
using i64 = long long;
constexpr int N = 1e5+5;

template <typename T>
struct Fenwick {
    int n;
    std::vector<T> a;

    Fenwick(int n_ = 0) {
        init(n_);
    }

    void init(int n_) {
        n = n_;
        a.assign(n, T{});
    }

    void add(int x, const T &v) {
        for (int i = x + 1; i <= n; i += i & -i) {
            a[i - 1] = a[i - 1] + v;
        }
    }

    T sum(int x) {
        T ans{};
        for (int i = x; i > 0; i -= i & -i) {
            ans = ans + a[i - 1];
        }
        return ans;
    }

    T rangeSum(int l, int r) {
        return sum(r) - sum(l);
    }

    int select(const T &k) {
        int x = 0;
        T cur{};
        for (int i = 1 << std::__lg(n); i; i /= 2) {
            if (x + i <= n && cur + a[x + i - 1] <= k) {
                x += i;
                cur = cur + a[x - 1];
            }
        }
        return x;
    }
};

int main() {
    std::cin.tie(nullptr)->sync_with_stdio(false);

    int n,m;
    std::cin >> n >> m;

    std::vector<i64> a(N,0),diff1(N,0),diff2(N,0);
    Fenwick<i64> fenwick1(N),fenwick2(N);
    for(int i = 1;i<=n;++i) {
        std::cin >> a[i];
        diff1[i] = a[i]-a[i-1];
        diff2[i] = i*diff1[i];
    }

    for(int i = 1;i<=n;++i) {
        fenwick1.add(i-1,diff1[i]);
        fenwick2.add(i-1,diff2[i]);
    }
    while(m--) {
        std::string op;int l,r;i64 d;
        std::cin >> op;
        if(op=="C") {
            std::cin >> l >> r >> d;
            fenwick1.add(l-1,d);
            fenwick2.add(l-1,l*d);
            fenwick1.add(r,-d);
            fenwick2.add(r,(r+1)*(-d));
        }else {//op=="Q"
            std::cin >> l >> r;
            i64 pre1 = fenwick1.sum(r)*(r+1)-fenwick2.sum(r);
            i64 pre2 = fenwick1.sum(l-1)*l-fenwick2.sum(l-1);
            std::cout << pre1-pre2 << '\n';
        }
    }
    
    return 0;
}

五、后记

等这几天回去再去看看线段树，后续可以类比一下线段树的写法和思路，线段树实际上更加全能，但是代码量真的太大，理解不透彻的人写着写着就Segment Fault了