W15 CF 763A

题意：给出一颗n个节点(1e5)树，每个点有一个颜色，判断是否存在一个点，以该节点为根，所有子树（不含该节点）为纯色（纯色：该子树所有点颜色相同）

sol:先假定以1为根，每个节点存两个值，g[x]表示以该节点为根，不含该节点的所有小子树是否都是纯色

f[x]表示加上该节点组成的一个大子树是否是纯色

首先判断g[1]值，如果不行，那么说明至少有一个小子树不是纯色，那么我们最后的根就要设定在这个子树里。

统计有多少个不是纯色的小子树或纯色但不和1颜色相同，如果>1,直接退。否则说明该节点所有祖先纯色，可以作为该点的一颗子树，上面合法，没有限制，向下找

#include<iostream>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#define N (100100)
using namespace std;
struct edge{
	int next,to;
	void Add(int Next,int T){
		next=Next; to=T;
	}
}q[4*N];
int C[N],h[N];
bool f[N],g[N];
void Dfs1(int x,int fa){
	int i,y;
	f[x]=g[x]=1;
	for (i=h[x];i;i=q[i].next){
		y=q[i].to; 
		if (y==fa) continue;
		Dfs1(y,x);
		if (!f[y]) f[x]=g[x]=0;
		if (C[x]!=C[y]) f[x]=0;
	}
}
void Dfs2(int x,int fa){
	if (g[x]){
		printf("YES\n%d",x); exit(0);
	}
	int i,y,num=0,pos=0;
	for (i=h[x];i;i=q[i].next){
		y=q[i].to; 
		if (y==fa) continue;
		if (!f[y]||C[y]!=C[x]){
			num++; pos=y;
		}
	}
	if (num>1){
		printf("NO"); exit(0);
	}
	if (C[pos]!=C[x])
		if (g[pos])
			printf("YES\n%d",pos); 
		else printf("NO\n");
	else
		Dfs2(pos,x);
}
int main(){
	int i,n; scanf("%d",&n);
	for (i=1;i<n;i++){
		int x,y; scanf("%d %d",&x,&y);
		q[2*i-1].Add(h[x],y); h[x]=2*i-1; 
		q[2*i].Add(h[y],x); h[y]=2*i; 
	}
	for (i=1;i<=n;i++) scanf("%d",&C[i]);
	Dfs1(1,0);
	Dfs2(1,0);
}