Codeforces Round #395 (Div.2) C-Timofey and a tree 树型dp

题目链接：http://codeforces.com/contest/764/problem/C

一道典型的树形dp题目，但是由于很久没有写代码以及模板也短时间找不到。。于是花了很长时间才写出来，差一点就没写出来orz，于是这里来记一下。

题目大意是说给一颗树，树的每一个节点都有颜色，问是否存在一个节点，使得把这一个节点当作根，其所有子树上的颜色是一致的。

随便选择一个节点做根，将无根树转为有根树，进行两遍dp。

第一遍dp从叶子节点推到根节点，计算以该节点为根的子树的颜色，如果有不同的颜色则记为-1

第二遍dp从根节点推到叶子节点，根据第一遍的dp值来计算某个节点向上的方向上的颜色，有不同颜色则记为-1。

这个时候答案就很明确了，找到一个节点，它的儿子节点为根的子树颜色统一（即所有儿子节点的第一遍dp值不为-1），并且向上的所以节点颜色统一（即第二遍的dp值不为-1）

时间复杂度为O(n)

辣鸡代码如下：

#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <map>
#include <queue>
#include <vector>
using namespace std;

const int max_n=300000;

struct tree
{
	int fa;
	int col;	//子树的dp值 
	int up;		//向上走的dp值 
	vector<int> son;
} Tree[max_n];

vector<int> G[max_n];
int color[max_n];	//节点值 

int search_up(int x)	//从叶子节点向上第一遍dp，计算子树dp值 
{
	int m=Tree[x].son.size();
	int mark=true;
	for(int i=0;i<m;i++)
	{
		int v=Tree[x].son[i];
		int c=search_up(v);
		if(c==-1 || c!=color[x]) mark=false;
	}
	
	return Tree[x].col=mark?color[x]:-1;
}

void search_down(int x,int p)		//从根节点向下递归计算向上走的dp值 
{
	if(p==-1)
	{
		Tree[x].up=-1;
		int m=Tree[x].son.size();
		for(int i=0;i<m;i++)	//向下递归 
			search_down(Tree[x].son[i],-1);
		return ;
	}
	
	int c;
	int fa=Tree[x].fa;
	if(fa==-1) Tree[x].up=c=color[x];
	else
	{
		
		c=color[fa];
		if(c!=Tree[fa].up) Tree[x].up=c=-1;
		int m=Tree[fa].son.size();
		for(int i=0;i<m;i++)	//枚举目标子树的父亲节点除了目标节点之外的所有子数，计算目标子树向上的dp值 
		{
			if(c==-1) break;
			int v=Tree[fa].son[i];
			if(v==x) continue;
			if(c!=Tree[v].col)
				c=-1;
		}
		Tree[x].up=c;
	}
	
	int m=Tree[x].son.size();
	for(int i=0;i<m;i++)	//向下递归 
		search_down(Tree[x].son[i],c);
}

int main()
{
	int n;
	cin>>n;
	for(int i=0;i<n-1;i++)	//建无根树 
	{
		int u,v;
		scanf("%d%d",&u,&v);
		u--;v--;
		G[u].push_back(v);
		G[v].push_back(u);
	}
	for(int i=0;i<n;i++)	scanf("%d",&color[i]);
	
	Tree[0].fa=-1; 
	queue<int> u;
	u.push(0);
	while(!u.empty())	//以0为根节点无根树转有根树 
	{
		int v=u.front();
		u.pop();
		int m=G[v].size();
		for(int i=0;i<m;i++)
		{
			int b=G[v][i];
			if(Tree[v].fa==b) continue;
			Tree[b].fa=v;
			Tree[v].son.push_back(b);
			u.push(b);
		}
	}
	
	search_down(0);	//递归计算子树dp值 
	search_up(0,1);	//递归计算向上走的dp值 
	
	for(int i=0;i<n;i++)	//根据dp值计算答案 
	{
		if(Tree[i].up==-1) continue;
		bool mark=true;
		int m=Tree[i].son.size();
		for(int j=0;j<m;j++)
		{
			if(Tree[Tree[i].son[j]].col==-1)
			{
				mark=false;
				break;
			}
		}
		if(mark)
		{
			cout<<"YES"<<endl;
			cout<<i+1<<endl;
			return 0;
		}
	}
	
	cout<<"NO"<<endl;
	return 0;
	
}