\(O(n^2)\)显然不过我们应该优化成\(O(nlogn)\)
采用树上启发式合并
仿照树链剖分的思想,对于每一个位置,我们先处理所有的轻儿子,然后处理重儿子,统计当前节点的答案,最后把轻儿子删掉就可以了。
这样全局一个桶就够用了。
#include<iostream>
#include<cstdio>
#include<algorithm>
#include<cstring>
using namespace std;
int n;
int x,y;
int col[1000005];
struct e{
int to;
int ne;
}ed[3000005];
int p;
int head[1000005];
void add(int f,int to){
p++;
ed[p].to=to;
ed[p].ne=head[f];
head[f]=p;
return ;
}
int siz[2000005];
int ns;
int ans[2000005];
int tem;
int son[1000005];
int exi[1000005];
void dfs(int no,int fa){
siz[no]++;
for(int i=head[no];i;i=ed[i].ne){
if(ed[i].to==fa) continue;
dfs(ed[i].to,no);
siz[no]+=siz[ed[i].to];
if(siz[ed[i].to]>siz[son[no]]){
son[no]=ed[i].to;
}
}
}
void work(int no,int fa,int ke){
if(!exi[col[no]]){
tem++;
}
exi[col[no]]+=ke;
for(int i=head[no];i;i=ed[i].ne){
int v=ed[i].to;
if(v==fa||v==ns) continue;
work(v,no,ke);
}
}
void dsu(int no,int fa,int f){
for(int i=head[no];i;i=ed[i].ne){
if(ed[i].to==fa||ed[i].to==son[no]) continue;
dsu(ed[i].to,no,0);
}
if(son[no]) {
dsu(son[no],no,1);
ns=son[no];
}
work(no,fa,1);
ns=0;
ans[no]=tem;
if(!f){
work(no,fa,-1);
tem=0;
}
}
int main(){
scanf("%d",&n);
for(int i=1;i<n;++i){
scanf("%d%d",&x,&y);
add(x,y);
add(y,x);
}
for(int i=1;i<=n;++i){
scanf("%d",&col[i]);
}
dfs(1,0);
//cout<<2344;
dsu(1,0,1);
scanf("%d",&n);
for(int i=1;i<=n;++i){
scanf("%d",&x);
printf("%d\n",ans[x]);
}
return 0;
}