树链剖分 - Acwing2568

树链剖分 - Acwing2568

树链剖分：一个强行增加代码量的树形结构预处理，将树拆分成若干条链，以便用线段树等其他数据结构在树上进行区间操作。套上树链剖分后，将使得原本的各项区间操作复杂度再乘上一个logn(最坏情况下)。

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;

const int N = 1e5+10;

vector<int> link[N];

int n, m;
int a, b, op, k;
int dfs_id[N];      // 节点的dfs序id
int head_id[N];     // 节点所在重链的起始节点的id
int fa[N];          // 父节点id
int subtree_start[N];       // 子树对应的dfs序左端点
int subtree_end[N];         // 子树对应的dfs序右端点
int dep[N];         // 节点深度
int heavy_son[N];   // 重儿子
int is_heavy[N];    // 该节点是否为重儿子

ll tree[N<<2];
ll lazy[N<<2];

ll arr[N];
ll new_arr[N];


void push_up(int rt){
    tree[rt] = tree[rt<<1] + tree[rt<<1|1];
}

void push_down(int rt,int l,int r){
    if(lazy[rt]){
        int mid = (l+r)>>1;
        tree[rt<<1] += (mid-l+1) * lazy[rt];
        tree[rt<<1|1] += (r-mid) * lazy[rt];
        lazy[rt<<1] += lazy[rt];
        lazy[rt<<1|1] += lazy[rt];
        lazy[rt] = 0;
    }
}

void build(int rt,int l,int r){
    if(l == r){
        tree[rt] = new_arr[l];
    }else{
        int mid = (l+r)>>1;
        build(rt<<1, l, mid);
        build(rt<<1|1, mid+1, r);
        push_up(rt);
    }
}

void modify(int rt,int l,int r,int ml,int mr,int v){
    if(ml <= l && mr >= r){
        tree[rt] += v * (r-l+1);
        lazy[rt] += v;
    }else{
        push_down(rt, l, r);
        int mid = (l+r)>>1;
        if(ml <= mid){
            modify(rt<<1, l, mid, ml, mr, v);
        }
        if(mr > mid){
            modify(rt<<1|1, mid+1, r, ml, mr, v);
        }
        push_up(rt);
    }
}

ll query(int rt,int l,int r,int ql,int qr){
    if(ql <= l && qr >= r){
        return tree[rt];
    }else{
        push_down(rt, l, r);
        ll ans = 0;
        int mid = (l+r)>>1;
        if(ql <= mid){
            ans += query(rt<<1, l, mid, ql, qr);
        }
        if(qr > mid){
            ans += query(rt<<1|1, mid+1, r, ql, qr);
        }
        return ans;
    }
}


void add_edge(int a,int b){
    link[a].push_back(b);
    link[b].push_back(a);
}

int dfs1(int cur,int prev,int d){
    fa[cur] = prev;     // 记录父节点
    dep[cur] = d;       // 记录深度
    int tot_size = 1;   // 以当前节点为根的子树的大小
    int max_size = -1,max_root = 0, cur_size = 0;
    
    for(auto it = link[cur].begin(); it != link[cur].end(); ++it){ // 遍历其每一个儿子
        if(*it != prev){ // 排除父节点
            cur_size = dfs1(*it, cur, d+1);
            if(cur_size > max_size){ // 更新重儿子
                max_size = cur_size;
                max_root = *it;
            }
            tot_size += cur_size;
        }
    }
    heavy_son[cur] = max_root;
    return tot_size;
}


void dfs2(int cur,int h,int prev,int &id){
    head_id[cur] = h;   // 节点所在重链的起始节点的id
    dfs_id[cur] = id++; // 记录dfs序
    
    
    if(heavy_son[cur]){ // 如果当前节点存在重儿子
        dfs2(heavy_son[cur], h, cur, id);
    }
    
    for(auto it = link[cur].begin(); it != link[cur].end(); ++it){
        if(*it != prev && *it != heavy_son[cur]){ // 排除父节点 & 排除重儿子(重儿子如果存在则已经被优先遍历过)
                dfs2(*it, *it, cur, id);
        }
    }

    subtree_start[cur] = dfs_id[cur];
    subtree_end[cur] = id-1;
}

void build_lct(){
    // dfs1(int cur,int prev,int d)
    dfs1(1, 0, 1);
    int id = 1;
    // dfs2(int cur,int h,int prev,int &id)
    dfs2(1, 1, 0, id);
}

void modify_lct(int a,int b,int v){ // 在a和b的路径上增加v
    int ha,hb;
    while(1){
        ha = head_id[a];
        hb = head_id[b];
        
        if(ha == hb){
            a = dfs_id[a];
            b = dfs_id[b];
            if(a > b) swap(a, b);
            modify(1, 1, n, a, b, v);
            break;
        }else{
            if(dep[ha] < dep[hb]){
                swap(a, b);
                swap(ha, hb);
            }
            modify(1, 1, n, dfs_id[ha], dfs_id[a], v); // 这是一条重链
            a = fa[ha];
        }
    }
}


ll query_lct(int a,int b){
    int ha, hb;
    ll ans = 0;
    while (1) {
        ha = head_id[a];
        hb = head_id[b];
        
        if(ha == hb){
            a = dfs_id[a];
            b = dfs_id[b];
            if(a > b) swap(a, b);
            ans += query(1, 1, n, a, b);
            break;
        }else{
            if(dep[ha] < dep[hb]){
                swap(a, b);
                swap(ha, hb);
            }
            ans += query(1, 1, n, dfs_id[ha], dfs_id[a]);
            a = fa[ha];
        }
    }
    return ans;
}


void modify_sub(int id,int v){
    modify(1, 1, n, subtree_start[id], subtree_end[id], v);
}

ll query_sub(int id){
    return query(1, 1, n, subtree_start[id], subtree_end[id]);
}

int main(){
    scanf("%d",&n);
    for(int i = 1; i <= n; ++i){
        scanf("%d",&arr[i]);
    }
    
    for(int i = 0; i < n-1; ++i){
        scanf("%d%d",&a,&b);
        add_edge(a, b);
    }
    
    build_lct();
    
    for(int i = 1; i <= n; ++i){
        new_arr[dfs_id[i]] = arr[i];
    }
    
    build(1, 1, n);
    
    scanf("%d",&m);
    while (m--) {
        scanf("%d",&op);
        if(op == 1){
            scanf("%d%d%d",&a,&b,&k);
            modify_lct(a, b, k);
        }else if(op == 2){
            scanf("%d%d",&a,&k);
            modify_sub(a, k);
        }else if(op == 3){
            scanf("%d%d",&a,&b);
            printf("%lld\n",query_lct(a, b));
        }else{
            scanf("%d",&a);
            printf("%lld\n",query_sub(a));
        }
    }
    
    return 0;
}