题目链接：http://codeforces.com/contest/764/problem/C
题意：给出一个树，然后各个节点有对应的颜色，问是否存在以一个点为根节点子树的颜色都一样。
这里的子树颜色都一样表示分出来的子树各自颜色一样。
就是随意找一个树叶节点然后一直遍历，直到找到一个和他颜色不一样的节点，那么要么取不一样的节点
为根或者取最后一个一样的节点为根。然后就dfs一遍。
#include <iostream>
#include <cstring>
#include <vector>
using namespace std;
const int M = 1e5 + 10;
int c[M] , ans , ans2;
vector<int>vc[M];
void dfs1(int sta , int pre , int co) {
if(c[sta] != co) {
ans = sta;
ans2 = pre;
return;
}
int len = vc[sta].size();
for(int i = 0 ; i < len ; i++) {
int v = vc[sta][i];
if(v != pre) {
dfs1(v , sta , co);
}
}
}
bool dfs(int sta , int pre , int co) {
if(c[sta] != co) {
return false;
}
int len = vc[sta].size();
int gg = true;
for(int i = 0 ; i < len ; i++) {
int v = vc[sta][i];
if(v != pre) {
if(!dfs(v , sta , co)) {
gg = false;
}
}
}
if(!gg) {
return false;
}
else {
return true;
}
}
int main() {
int n , u , v;
cin >> n;
for(int i = 0 ; i < n - 1 ; i++) {
cin >> u >> v;
vc[u].push_back(v);
vc[v].push_back(u);
}
for(int i = 1 ; i <= n ; i++) {
cin >> c[i];
}
int sta = 0;
for(int i = 1 ; i <= n ; i++) {
int len = vc[i].size();
if(len == 1) {
sta = i;
break;
}
}
ans = -1 , ans2 = -1;
dfs1(sta , sta , c[sta]);
if(ans == -1) {
cout << "YES" << endl;
cout << sta << endl;
}
else {
int len = vc[ans].size();
bool flag = true;
for(int i = 0 ; i < len ; i++) {
int v = vc[ans][i];
flag = dfs(v , ans , c[v]);
if(!flag)
break;
}
int flag1 = true;
len = vc[ans2].size();
for(int i = 0 ; i < len ; i++) {
int v = vc[ans2][i];
flag1 = dfs(v , ans2 , c[v]);
if(!flag1)
break;
}
if(flag || flag1) {
cout << "YES" << endl;
if(flag) {
cout << ans << endl;
return 0;
}
if(flag1) {
cout << ans2 << endl;
}
}
else {
cout << "NO" << endl;
}
}
return 0;
}