洛谷U41492 樹上數顏色 樹上啓發式合併

題目鏈接：https://www.luogu.com.cn/problem/U41492
題意：給一棵根爲1的樹，每次詢問子樹顏色種類數
思路：樹上啓發式合併，每次計算時如果是第一次出現就計算貢獻，輕鏈的話要消除貢獻 

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
#define fi first
#define se second
#define ls rt << 1
#define rs rt << 1|1
#define po pop_back
#define pb push_back
#define mk make_pair
#define lson l, mid, ls
#define rson mid + 1, r, rs
#define pll pair<ll, ll>
#define pii pair<int, int>
#define ull unsigned long long
#define pdd pair<double, double>
const int mod = 1e9 + 9;
const int maxn = 2e5 + 10;
const int inf = 0x3f3f3f3f;
const ll linf = 0x3f3f3f3f3f3f3f3f;
int c[maxn], num[maxn], sz[maxn], son[maxn], sum[maxn], cnt[maxn], ans[maxn], Max, Son;
vector<int> e[maxn];

void dfs(int u, int fa) //統計重兒子信息
{
    sz[u] = 1;
    //sum[u] = num[u];
    for(auto v : e[u])
    {
        if(v == fa) continue;
        dfs(v, u);
        sz[u] += sz[v];
        //sum[u] += sum[v];
        if(sz[son[u]] < sz[v]) son[u] = v;
    }
}
void calc(int u, int fa, int val)
{
    if(cnt[c[u]] == 0) ++Max;
    cnt[c[u]] += val;
    for(auto v : e[u])
    {
        if(v == fa || v == Son) continue;
        calc(v, u, val);
    }
}
void dfs2(int u, int fa, int op) // 0輕鏈 1重鏈
{
    for(auto v : e[u])
    {
        if(v == fa || v == son[u]) continue;
        dfs2(v, u, 0);
    }
    if(son[u]) //重兒子計算貢獻 不消除
         dfs2(son[u], u, 1), Son = son[u];
    calc(u, fa, 1), Son = 0;
    ans[u] = Max;
    if(!op) calc(u, fa, -1), Max = 0; //消除貢獻
}

int main()
{
    int n;
    scanf("%d", &n);
    for(int i = 1; i < n; ++i)
    {
        int u, v;
        scanf("%d%d", &u, &v);
        e[u].pb(v), e[v].pb(u);
    }
    for(int i = 1; i <= n; ++i) scanf("%d", &c[i]);
    dfs(1, 0);
    dfs2(1, 0, 0);
    int m;
    scanf("%d", &m);
    while(m--)
    {
        int x;
        scanf("%d", &x);
        printf("%d\n", ans[x]);
    }
	return 0;
}