Codeforces 763A. Timofey and a tree

A. Timofey and a tree

题意：给一棵树，要求判断是否存在一个点，删除这个点后，所有连通块内颜色一样。$N,C \le 10^5$

想法：这个叫换根吧。先求出一个点合法即其儿子的子树内颜色一样，非该点子树的点颜色都一样。可以用DFS序解决。

#include< cstdio >

typedef long long ll;
template
inline void read(T&x)
{
	x=0;bool f=0;char c=getchar();
	while((c<'0'||c>'9')&&c!='-') c=getchar();if(c=='-')f=1, c=getchar();
	while(c>='0'&&c<='9'){x=x*10+c-'0'; c=getchar();}
	x=f?-x:x;
}
const int MAXN(100010);
struct Node{int nd,nx;}bot[MAXN<<1];int tot,first[MAXN];
void add(int a,int b){bot[++tot]=(Node){b,first[a]};first[a]=tot;}
int n,a,b,col[MAXN],Eu[MAXN],size[MAXN],num[MAXN],last[MAXN],cnt,Ans;
bool dif[MAXN];
void DFS(int x,int f)
{
	Eu[x]=++cnt; num[cnt]=x;size[x]=1;
	for(int v=first[x];v;v=bot[v].nx)
	if(bot[v].nd!=f)
	{
		DFS(bot[v].nd,x); size[x]+=size[bot[v].nd];
		dif[x]|=dif[bot[v].nd]|(col[x]!=col[bot[v].nd]);
	}
}
void DP(int x,int f)
{
	bool bf=false;
	for(int v=first[x];v;v=bot[v].nx)
	if(bot[v].nd!=f)
	{
		DP(bot[v].nd,x);
		bf|=dif[bot[v].nd];
	}
	if(!bf)
	{
//		fprintf(stderr,"%d\n",x);
		int L=Eu[x]-1,R=Eu[x]+size[x];
//		fprintf(stderr,"%d %d\n",L,R);
		if((!L||last[1]>=L)&&(R>n||last[R]>=n)&&(!L||R>n||col[num[1]]==col[num[n]]))Ans=x;
	}
}
int main()
{
//	freopen("C.in","r",stdin);
	read(n);
	for(int i=1;i<n;i++)
	{
		read(a);read(b);
		add(a,b);add(b,a);
	}
	for(int i=1;i<=n;i++)read(col[i]);
	DFS(1,0);
	for(int i=n;i>=1;i--)last[i]=(col[num[i]]==col[num[i+1]])?last[i+1]:i;
	DP(1,0);
	if(!Ans)printf("NO");
	else printf("YES\n%d\n",Ans);
	return 0;
}