Aragorn's Story - 树链剖分 - HDU 3966

Aragorn's Story - 树链剖分 - HDU 3966

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

typedef long long ll;

const int N = 5e4+5;
int n, m, p;

char op[5];
int a, b, c;

vector<int> edges[N];

int head[N];
int fa[N];
int dep[N];
int dfs_id[N];
int heavy_son[N];
int w[N];
int dfs_w[N];

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

int dfs1(int cur, int prev, int depth){
    fa[cur] = prev;
    dep[cur] = depth;
    int tot_size = 1, cur_size = 0, max_size = 0, max_root = 0;
    for(vector<int> :: iterator it = edges[cur].begin(); it != edges[cur].end(); ++it){
        if(*it == prev) continue;
        cur_size = dfs1(*it, cur, depth+1);
        tot_size += cur_size;
        if(cur_size > max_size){
            max_size = cur_size;
            max_root = *it;
        }
    }
    heavy_son[cur] = max_root;
    return tot_size;
}

void dfs2(int cur, int prev, int head_node, int& id){
    dfs_id[cur] = id++;
    dfs_w[dfs_id[cur]] = w[cur];
    
    head[cur] = head_node;
    if(heavy_son[cur]){
        dfs2(heavy_son[cur], cur, head_node, id);
    }
    for(vector<int> :: iterator it = edges[cur].begin(); it != edges[cur].end(); ++it){
        if(*it == prev || *it == heavy_son[cur]) continue;
        dfs2(*it, cur, *it, id);
    }
}

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 modify(int rt, int l, int r, int ml, int mr, int v){
    if(ml <= l && mr >= r){
        tree[rt] += (r-l+1) * v;
        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 modify_lct(int a, int b, int v){
    int ha, hb;
    while(1){
        ha = head[a];
        hb = head[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];
        }
    }
}

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


int main(){
    while(scanf("%d%d%d", &n, &m, &p)!=EOF){
        for(int i = 1; i <= n; ++i){
            scanf("%d", &w[i]);
            edges[i].clear();
            heavy_son[i] = 0;
        }

        for(int i = 1; i <= m; ++i){
            scanf("%d%d", &a, &b);
            edges[a].push_back(b);
            edges[b].push_back(a);
        }
        dfs1(1, 0, 1);
        int id = 1;
        dfs2(1, 0, 1, id);


        build(1, 1, n);

        while(p--){
            scanf("%s", op);
            if(*op == 'I'){
                scanf("%d%d%d", &a, &b, &c);
                modify_lct(a, b, c);
            }else if(*op == 'D'){
                scanf("%d%d%d", &a, &b, &c);
                modify_lct(a, b, -c);
            }else{
                scanf("%d", &a);
                printf("%lld\n", query(1, 1, n, dfs_id[a], dfs_id[a]));
            }
        }
    }
    return 0;
}