【树形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
,next[N<<1],en,col
,f
,fa
;
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;
}