U41492 树上数颜色(dsu on tree)

blackpink

\(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;
}