Colorful Tree

Colorful Tree

rk_no巨巨TQL

/*
  Author : lifehappy
*/
#include <bits/stdc++.h>

using namespace std;

const int N = 1e5 + 10;

int head[N], to[N << 1], nex[N << 1], cnt = 1;

int son[N], dep[N], fa[N], sz[N], rk[N], id[N], top[N], tot;

int n, m, value[N], ans[N];

void Add(int x, int y) {
    to[cnt] = y;
    nex[cnt] = head[x];
    head[x] = cnt++;
}

void dfs1(int rt, int f) {
    dep[rt] = dep[f] + 1;
    sz[rt] = 1, rk[++tot] = rt, id[rt] = tot;
    fa[rt] = f;
    for(int i = head[rt]; i; i = nex[i]) {
        if(to[i] == f) continue;
        dfs1(to[i], rt);
        sz[rt] += sz[to[i]];
        if(!son[rt] || sz[to[i]] > sz[son[rt]]) son[rt] = to[i];
    }
}

void dfs2(int rt, int tp) {
    top[rt] = tp;
    if(!son[rt]) return ;
    dfs2(son[rt], tp);
    for(int i = head[rt]; i; i = nex[i]) {
        if(to[i] == fa[rt] || to[i] == son[rt]) continue;
        dfs2(to[i], to[i]);
    }
}

int lca(int x, int y) {
    while(top[x] != top[y]) {
        if(dep[top[x]] < dep[top[y]]) swap(x, y);
        x = fa[top[x]];
    }
    return dep[x] < dep[y] ? x : y;
}

int dis(int x, int y) {
    x = rk[x], y = rk[y];
    return dep[x] + dep[y] - 2 * dep[lca(x, y)];
}

set<int> st[N];

void add(int x, int value) {
    if(st[value].size() == 0) {
        ans[value] = 0;
        st[value].insert(x);
        return ;
    }
    auto it = st[value].lower_bound(x);
    if(it == st[value].begin() || it == st[value].end()) {
        auto y = st[value].begin();
        auto z = st[value].rbegin();
        ans[value] += (dis(x, *y) + dis(x, *z) - dis(*y, *z)) >> 1;
    }
    else {
        auto y = it, z = it;
        y--;
        ans[value] += (dis(x, *y)+dis(x, *z)-dis(*y, *z))/2;
    }
    st[value].insert(x);
}

void del(int x, int value) {
    if (st[value].size() == 1){
        st[value].erase(x); ans[value] = -1;
        return;
    }
    st[value].erase(x);
    auto it = st[value].lower_bound(x);
    if (it == st[value].begin() || it == st[value].end()){
        auto y = st[value].begin();
        auto z = st[value].rbegin();
        ans[value] -= (dis(x, *y)+dis(x, *z)-dis(*y, *z))/2;
    }else{
        auto y = it, z = it; y--;
        ans[value] -= (dis(x, *y)+dis(x, *z)-dis(*y, *z))/2;
    }
}

int main() {
    // freopen("in.txt", "r", stdin);
    // freopen("out.txt", "w", stdout);
    // ios::sync_with_stdio(false), cin.tie(0), cout.tie(0);
    scanf("%d", &n);
    for(int i = 1; i < n; i++) {
        int x, y;
        scanf("%d %d", &x, &y);
        Add(x, y);
        Add(y, x);
    }
    dfs1(1, 0);
    dfs2(1, 1);
    memset(ans, -1, sizeof ans);
    for(int i = 1; i <= n; i++) {
        scanf("%d", &value[i]);
        add(id[i], value[i]);
    }
    scanf("%d", &m);
    for(int i = 1; i <= m; i++) {
        char op;
        cin >> op;
        if(op == 'Q') {
            int value;
            scanf("%d", &value);
            printf("%d\n", ans[value]);
        }
        else {
            int x, color;
            scanf("%d %d", &x, &color);
            del(id[x], value[x]);
            value[x] = color;
            add(id[x], value[x]);
        }
    }
    return 0;
}