牛客练习赛47 DongDong数颜色 树上启发式合并

题目链接：https://ac.nowcoder.com/acm/contest/904/E
题目大意：

思路：裸的树上轻重链启发式合并

#include <bits/stdc++.h>
#define LL long long
using namespace std;

int a[100005];//dfs序的颜色
int b[100005];//i节点的颜色
int c[100005];//颜色桶
int dfn[100005];//dfs序
int s[100005];//子树节点个数
int son[100005];//重链节点
int ans[100005];
int now=0, T=0;
vector<int> G[100005];

void dfs(int u, int fa){//得到dfs序，和重链的节点
    s[u]=1; dfn[u]=++T;
    for(auto x: G[u]){
        if(x!=fa){
            dfs(x, u); s[u]+=s[x];
            son[u]=s[x]>s[son[u]]?x:son[u];
        }
    }
}

void add(int u){//计算贡献
    now+=!c[u]++;
}

void DFS(int u, int fa, int k){//当前节点， 父亲节点， 是否保留c数组
    for(auto x: G[u]){//把除了重链的所有的子树的答案计算出来
        if(x!=fa&&x!=son[u]){
            DFS(x, u, 0);
        }
    }
    //计算u这棵树的答案
    if(son[u]) DFS(son[u], u, 1);//先计算重链
    for(auto x: G[u]){//轻链上的贡献
        if(x!=fa&&x!=son[u]){
            for(int i=0; i<s[x]; i++){
                add(a[dfn[x]+i]);//子树的dfs序是连续的
            }
        }
    }
    add(a[dfn[u]]); ans[u]=now;//把u自己加进去

    if(!k){//是否需要清空
        for(int i=now=0; i<s[u]; i++){
            c[a[dfn[u]+i]]=0;
        }
    }
}

int main(){
    int n, m, x, y, q; scanf("%d%d", &n, &m);
    for(int i=1; i<=n; i++){
        scanf("%d", &b[i]);
    }
    for(int i=1; i<n; i++){
        scanf("%d%d", &x, &y);
        G[x].push_back(y);
        G[y].push_back(x);
    }
    dfs(1, 0);
    for(int i=1; i<=n; i++){
        a[dfn[i]]=b[i];
    }
    DFS(1, 0, 1);
    while(m--){
        scanf("%d", &q);
        printf("%d\n", ans[q]);
    }

    return 0;
}