题目链接: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; }