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;
	
}