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