P2590 [ZJOI2008]树的统计

Posted 2021-01-11 wsmrxc

久违的树剖~

#include <vector>
#include <cstdio>
#include <cstring>
#include <cassert>
#include <iostream>
#include <algorithm>
using namespace std;
const int MAXN = 3e4 + 10;
inline int read(){
    char ch = getchar(); int x = 0; bool flag = false;
    while(!isdigit(ch)) flag |= (ch == '-'), ch = getchar();
    while(isdigit(ch)) x = x * 10 + ch - '0', ch = getchar();
    return flag ? -x : x;
}

int N, val[MAXN];
vector<int> g[MAXN];

int fa[MAXN], dep[MAXN], siz[MAXN], son[MAXN];
int dfs1(int u, int f, int depth){
    dep[u] = depth, fa[u] = f, siz[u] = 1;

    int tmp = -1;
    for(int i = 0; i < (int) g[u].size(); i++){
        int &v = g[u][i];
        if(v == f) continue;

        siz[u] += dfs1(v, u, depth + 1);
        if(tmp == -1 || siz[v] > siz[tmp]) tmp = v;
    }
    son[u] = tmp; return siz[u];
}

int top[MAXN], idx[MAXN], reidx[MAXN];
void dfs2(int u, int t){
    static int cntnode = 0;
    idx[u] = ++cntnode, top[u] = t, reidx[cntnode] = u;

    if(son[u] == -1) return;

    dfs2(son[u], t);
    for(int i = 0; i < (int) g[u].size(); i++){
        int &v = g[u][i];
        if(v == fa[u] || v == son[u]) continue;
        dfs2(v, v);
    }
}

namespace stree
{
    #define mid ((l + r) >> 1)
    #define ls (o << 1)
    #define rs ((o << 1) | 1)

    struct Node
    {
        int sum, val;
    }node[MAXN << 2];

    inline void pushup(int o){
        node[o].sum = node[ls].sum + node[rs].sum,
        node[o].val = max(node[ls].val, node[rs].val);
    }
                         
    void build(int o, int l, int r){
        if(l == r) return node[o].val = node[o].sum = val[reidx[l]], void();
        build(ls, l, mid), build(rs, mid + 1, r);
        return pushup(o);
    }

    void modify(int o, int l, int r, int p, int v){
        if(l == r) return node[o].val = node[o].sum = v, void();
        if(p <= mid) modify(ls, l, mid, p, v);
        else modify(rs, mid + 1, r, p, v);
        return pushup(o);
    }

    int query(int o, int l, int r, int a, int b, bool type) { //true -> max, false -> sum
        if(l > b || r < a) return type ? -(1 << 30) : 0;
        if(a <= l && r <= b) return type ? node[o].val : node[o].sum;
        if(type) return max(query(ls, l, mid, a, b, true), query(rs, mid + 1, r, a, b, true));
        else return query(ls, l, mid, a, b, false) + query(rs, mid + 1, r, a, b, false);
    }
}

int main(){
    // freopen("p2590.in", "r", stdin);
    // freopen("p2590.out", "w", stdout); 
    cin>>N;
    for(int i = 1; i < N; i++){
        int u = read(), v = read();
        g[u].push_back(v), g[v].push_back(u);
    }
    for(int i = 1; i <= N; i++) val[i] = read();
    dfs1(1, 0, 1), dfs2(1, 1);
    memset(stree::node, -0x3f, sizeof(stree::node));
    stree::build(1, 1, N);
    int Q; cin>>Q; char ch[10];
    while(Q--){
        using namespace stree;
        scanf(" %[A-Z]", ch); int u = read(), v = read();
        
        if(ch[0] == 'C')
            modify(1, 1, N, idx[u], v);
        else {
            if(ch[1] == 'M'){
                int ans = -(1 << 30);
                while(top[u] != top[v]) {
                    if(dep[top[u]] < dep[top[v]]) swap(u, v);
                    ans = max(ans, query(1, 1, N, idx[top[u]], idx[u], true));
                    u = fa[top[u]];
                }
                if(dep[u] < dep[v]) swap(u, v);
                ans = max(ans, query(1, 1, N, idx[v], idx[u], true));
                printf("%d
", ans);
            }
            else {
                int ans = 0;
                while(top[u] != top[v]) {
                    if(dep[top[u]] < dep[top[v]]) swap(u, v);
                    ans += query(1, 1, N, idx[top[u]], idx[u], false);
                    u = fa[top[u]];
                }
                if(dep[u] < dep[v]) swap(u, v);
                ans += query(1, 1, N, idx[v], idx[u], false);
                printf("%d
", ans);
            }
        }
    }
    return 0;
}