dsu on tree 求调
  • 板块学术版
  • 楼主After_light
  • 当前回复2
  • 已保存回复2
  • 发布时间2024/8/24 09:07
  • 上次更新2024/8/24 09:16:46
查看原帖
dsu on tree 求调
554803
After_light楼主2024/8/24 09:07

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

2024/8/24 09:07
加载中...