树状数组和线段树的实现

如何推导？查看本人的知乎这篇文章：前缀和改进之树状数组 - 知乎 (zhihu.com)

树状数组

功能

  • 单点修改，区间求和 O(logN)

  • 可以在，O(logN)的时间内，动态求前缀和 ( 对比前缀和，修改的话需要O(N) )

注意：树状数组下标一定要从1开始

树状数组初始化

1
for(int i=1 ; i<=n ; i++) add(i,a[i]);

小小地解释一下原理：

  • 树状数组中奇数下标和原数组相等

  • c[x] = 原数组中下标从(x-lowbit(x)，x]的和 左开右闭

三个基本操作：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
int a[N],c[N];//a[N]为原数组，c[N]为树状数组

int lowbit(int x){
    return x & -x;
}

//一般情况下只支持add，如果要实现修改某一个数，可以add他们的差值
void add(int x , int v ){
    for(int i=x;i<=n;i+=lowbit(i)) c[i] += v;    
}

//查询
int query(int x){
    int res = 0;
    for(int i = x ; i>=1 ; i-=lowbit(i) )  res += c[i];//for循环截止条件写成 i!=0 即 i也可以
    return res;
}

线段树

每个结点都是一个结构体

1
2
3
4
struct Node{
    int L,R;
    int sum;
}tr[N*4];

线段树存储

注意结点个数要开 4n 其中n是有多少个数

类似于堆的存储方式，存在一维数组，对于一个序号为x的节点

  • 父节点为 x/2 (下取整) x>>1
  • 左子节点为 x*2 x<<1
  • 右子节点为 x*2+1 x<< 1 | 1

基本操作

  • 单点修改 O(logN)

  • 区间查询 O(logN)

四个函数

pushup 用子节点信息更新当前节点信息

1
2
3
void pushup(int u){
    tr[u].sum = tr[u<<1].sum + tr[u<<1 | 1].sum;
}

build 在一段区间上初始化线段树

1
2
3
4
5
6
7
8
9
void build(int u, int l, int r){
    if(l == r) tr[u] = {l,r,w[r]};
    else{
        tr[u] = {l,r};
        int mid = tr[u].l + tr[u].r >> 1;
        build(u<<1,l,mid), build(u<<1 | 1 , mid+1 ,r);
        pushup(u);
    }
}

modify 修改某个点的值

1
2
3
4
5
6
7
8
9
void modify(int u, int x, int v){
    if(tr[u].l == tr[u].r ) tr[u].sum += v;
    else{
        int mid = tr[u].l + tr[u].r >> 1;
        if(x <= mid) modify(u<<1, x , v);
        else modify(u<<1|1 , x , v);
        pushup(u);
    }
}

query 查询

1
2
3
4
5
6
7
8
9
10
int query(int u, int l, int r){
    if(tr[u].l>=l && tr[u].r<=r) return tr[u].sum;
    else{
        int mid = tr[u].l + tr[u].r >> 1;
        int sum = 0;
        if(l <= mid) sum += query(u<<1 , l ,mid);
        if(r > mid)  sum += query(u<<1|1,l,r);
        return sum
    }
}