【树形DP】Codeforces Round #395 (Div. 2) C. Timofey and a tree

标题写的树形DP是瞎扯的。

先把1看作根。

预处理出f[i]表示以i为根的子树是什么颜色，如果是杂色的话，就是0。

然后从根节点开始转移，转移到某个子节点时，如果其子节点都是纯色，并且它上面的那一坨结点也是纯色，就输出解。

否则如果其上面的一坨是纯色，并且其子节点有且只有一个杂色的时候，就递归处理该子节点。

#include<cstdio>
#include<cstdlib>
using namespace std;
#define N 100050
int v[N<<1],first[N],next[N<<1],en,col[N],f[N],fa[N];
void AddEdge(int U,int V)
{
	v[++en]=V;
	next[en]=first[U];
	first[U]=en;
}
void dfs(int U)
{
	bool ok=1;
	for(int i=first[U];i;i=next[i]) if(v[i]!=fa[U])
	  {
	  	fa[v[i]]=U;
	  	dfs(v[i]);
	  	if(f[v[i]]!=col[U])
	  	  ok=0;
	  }
	if(ok)
	  f[U]=col[U];
}
void df2(int U)
{
	bool Got=1;
	for(int i=first[U];i;i=next[i]) if(v[i]!=fa[U])
	  if(!f[v[i]])
	    {
	      Got=0;
	      break;
	    }
	if(Got)
	  {
	  	printf("YES\n%d\n",U);
	  	exit(0);
	  }
	int cnt=0,vi;
	for(int i=first[U];i;i=next[i]) if(v[i]!=fa[U])
	  if(!f[v[i]])
	    {
	      ++cnt;
	      vi=v[i];
	    }
	if(cnt>1)
	  return;
	if(col[U]!=col[1])
	  return;
	for(int i=first[U];i;i=next[i]) if(v[i]!=fa[U] && v[i]!=vi)
	  if(f[v[i]]!=col[1])
	    return;
	df2(vi);
}
int n;
int main()
{
	//freopen("c.in","r",stdin);
	int x,y;
	scanf("%d",&n);
	for(int i=1;i<n;++i)
	  {
	  	scanf("%d%d",&x,&y);
	  	AddEdge(x,y);
	  	AddEdge(y,x);
	  }
	for(int i=1;i<=n;++i)
	  scanf("%d",&col[i]);
	dfs(1);
	df2(1);
	puts("NO");
	return 0;
}