【BZOJ 1036】 [ZJOI2008]树的统计

题目来源：BZOJ 1036

思路：

练习树链剖分的例题，注意修改操作，所给的信息是树上的点的编号，不是线段树里的编号。

代码：

#include <cstdio>
#include <iostream>
#include <algorithm>
#define mid ((l+r)>>1)
#define lch (now<<1)
#define rch ((now<<1)+1)
const int maxn = 30010;
const int inf = 2e9;
int n, m, ct;
int val[maxn], v[maxn];
int head[maxn], next[maxn<<1], to[maxn<<1], cnt;
int fa[maxn], dep[maxn], size[maxn], son[maxn];
int pos[maxn], top[maxn];
int t_max[maxn<<2], t_sum[maxn<<2];
void add(int a, int b){
    next[++cnt] = head[a];
    head[a] = cnt, to[cnt] = b;
}
void dfs1(int now){
    int mx1 = 0, mx2 = 0;
    size[now] = 1;
    for(int i = head[now]; i; i = next[i]){
        if(to[i] == fa[now]) continue;
        fa[to[i]] = now, dep[to[i]] = dep[now] + 1;
        dfs1(to[i]);
        size[now] += size[to[i]];
        if(mx1 < size[to[i]]) mx1 = size[to[i]], mx2 = to[i];
    }
    son[now] = mx2;
}
void dfs2(int now){
    pos[now] = ++ ct;
    if(son[now] != 0) top[son[now]] = top[now], dfs2(son[now]);
    for(int i = head[now]; i; i = next[i]){
        if(to[i] == son[now] || to[i] == fa[now]) continue;
        top[to[i]] = to[i], dfs2(to[i]);
    }
}
void build(int now, int l, int r){
    if(l == r){t_max[now] = v[l], t_sum[now] = v[l]; return;}
    build(lch, l, mid);
    build(rch, mid+1, r);
    t_max[now] = std::max(t_max[lch], t_max[rch]);
    t_sum[now] = t_sum[lch] + t_sum[rch];
}
int que_max(int now, int l, int r, int pos1, int pos2){
    if(l == pos1 && r == pos2) return t_max[now];
    if(pos2 <= mid) return que_max(lch, l, mid, pos1, pos2);
    else if(pos1 >= mid+1) return que_max(rch, mid+1, r, pos1, pos2);
    else return std::max(que_max(lch, l, mid, pos1, mid), que_max(rch, mid+1, r, mid+1, pos2));
}
int que_sum(int now, int l, int r, int pos1, int pos2){
    if(l == pos1 && r == pos2) return t_sum[now];
    if(pos2 <= mid) return que_sum(lch, l, mid, pos1, pos2);
    else if(pos1 >= mid+1) return que_sum(rch, mid+1, r, pos1, pos2);
    else return que_sum(lch, l, mid, pos1, mid) + que_sum(rch, mid+1, r, mid+1, pos2);
}
void modify(int now, int l, int r, int pos, int k){
    if(l == r){t_max[now] = k, t_sum[now] = k; return;}
    if(pos <= mid) modify(lch, l, mid, pos, k);
    else modify(rch, mid+1, r, pos, k);
    t_max[now] = std::max(t_max[lch], t_max[rch]);
    t_sum[now] = t_sum[lch] + t_sum[rch];
}
int que_max(int x, int y){
    int res = -inf;
    while(top[x] != top[y]){
        if(dep[top[x]] < dep[top[y]]) std::swap(x, y);
        res = std::max(res, que_max(1, 1, n, pos[top[x]], pos[x]));
        x = fa[top[x]];
    }
    if(pos[x] > pos[y]) std::swap(x, y);
    res = std::max(res, que_max(1, 1, n, pos[x], pos[y]));
    return res;
}
int que_sum(int x, int y){
    int res = 0;
    while(top[x] != top[y]){
        if(dep[top[x]] < dep[top[y]]) std::swap(x, y);
        res += que_sum(1, 1, n, pos[top[x]], pos[x]);
        x = fa[top[x]];
    }
    if(pos[x] > pos[y]) std::swap(x, y);
    res += que_sum(1, 1, n, pos[x], pos[y]);
    return res;
}
int main(){
    scanf("%d", &n);
    for(int i = 1, a, b; i < n; i ++){scanf("%d%d", &a, &b), add(a, b), add(b, a);}
    for(int i = 1; i <= n; i ++) scanf("%d", &val[i]);
    scanf("%d", &m);
    dep[1] = 1, dfs1(1);
    top[1] = 1, dfs2(1);
    for(int i = 1; i <= n; i ++) v[pos[i]] = val[i];
    build(1, 1, n);
    for(int i = 1, a, b; i <= m; i ++){
        char ch[20];
        scanf("%s%d%d", ch, &a, &b);
        if(ch[1] == 'M') printf("%d\n", que_max(a, b));
        if(ch[1] == 'S') printf("%d\n", que_sum(a, b));
        if(ch[1] == 'H') modify(1, 1, n, pos[a], b);
    }
    return 0;
}