U41492 樹上數顏色(dsu on tree)

#include<bits/stdc++.h>
#define endl '\n'
#define PB push_back
#define ll long long
using namespace std;
const int maxn = 1e5 + 10;
int n, m, sum, hson;
int siz[maxn], son[maxn], num[maxn], c[maxn], ans[maxn];
vector<int> g[maxn];
void dfs(int u, int f){
    siz[u] = 1;
    for(auto v : g[u]){
        if(v == f) continue;
        dfs(v, u);
        siz[u] += siz[v];
        if(siz[v] > siz[son[u]]) son[u] = v;
    }
}
void cal(int u, int f, int k){
    if(!num[c[u]]) sum ++;
    num[c[u]] += k;
    for(auto v : g[u]){
        if(v == f || v == hson) continue;
        cal(v, u, k);
    }
}
void dsu(int u, int f, int h){
    for(auto v : g[u]){
        if(v == f || v == son[u]) continue;
        dsu(v, u, 0);
    }
    if(son[u]) dsu(son[u], u, 1), hson = son[u];
    cal(u, f, 1); hson = 0;
    ans[u] = sum;
    if(!h) cal(u, f, -1), sum = 0;
}
int main()
{
    scanf("%d", &n);
    int u, v, x;
    for(int i = 1 ; i < n ; ++ i){
        scanf("%d %d", &u, &v);
        g[u].PB(v); g[v].PB(u);
    }
    for(int i = 1 ; i <= n ; ++ i) scanf("%d", &c[i]);
    scanf("%d", &m);
    dfs(1, 0);
    dsu(1, 0, 1);
    while(m --){
        scanf("%d", &x);
        cout << ans[x] << endl;
    }

    return 0;
}