题目链接:https://www.luogu.com.cn/problem/U41492
code:
#include<bits/stdc++.h>
#define F(i,a,b) for(int i=a;i<=b;i++)
#define R(i,a,b) for(int i=a;i>=b;i--)
using namespace std;
const int N=1e6+7;
int n,col[N],dep[N],fa[N],sz[N],son[N],vis[N],sum,ans[N],dfn[N],idx,id[N];
struct edge{
int nxt,to;
edge(int _nxt=0,int _to=0){
nxt=_nxt,to=_to;
}
}e[N];
int head[N],cnt;
inline void add_edge(int from,int to){
e[++cnt]=edge(head[from],to);
head[from]=cnt;
}
inline void dfs1(int now,int f){
sz[now]=1;
dep[now]=dep[f]+1;
dfn[++idx]=now;
id[now]=idx;
fa[now]=f;
for(int i=head[now];i;i=e[i].nxt){
if(e[i].to==f)continue;
dfs1(e[i].to,now);
sz[now]+=sz[e[i].to];
if(sz[son[now]]<sz[e[i].to]){
son[now]=e[i].to;
}
}
}
inline void add(int x){
vis[col[x]]++;
if(vis[col[x]]==1) sum++;
}
inline void del(int x){
vis[col[x]]--;
if(!vis[col[x]]) sum--;
}
inline void dfs2(int now,int f,bool k){
for(int i=head[now];i;i=e[i].nxt){
if(e[i].to!=f&&e[i].to!=son[now]) dfs2(e[i].to,now,false);
}
if(son[now]) dfs2(son[now],now,true);
add(now);
for(int i=head[now];i;i=e[i].nxt){
if(e[i].to==f||e[i].to==son[now]) continue;
add(e[i].to);
F(j,id[e[i].to],id[e[i].to]+sz[e[i].to]-1){
add(dfn[j]);
}
}
ans[now]=sum;
if(!k){
F(i,id[now],id[now]+sz[now]-1){
del(dfn[i]);
}
}
}
int main(){
cin>>n;
F(i,1,n-1){
int from,to;
cin>>from>>to;
add_edge(from,to);
add_edge(to,from);
}
F(i,1,n) cin>>col[i];
dfs1(1,0);
dfs2(1,1,0);
int m;cin>>m;
F(i,1,m){
int x;
cin>>x;
cout<<ans[x]<<"\n";
}
return 0;
}