CodeForces 764C Timofey and a tree

题目链接：http://codeforces.com/contest/764/problem/C
题意：给你一棵树，每个节点的颜色为col[i]，让你从这棵树中任意选择一个节点为根节点，使得所有的子树，在子树内部的所有颜色都是一样的，如果存在这样一个点那就数输出YES和这点，如果答案有多个，则输出任意一个，否则输出NO
解析：假设有cnt条边的两个端点颜色不一样，所以这cnt条边肯定是全都连在根节点上的，那么我们只需要记录，端点颜色不同的边的每个端点出现了多少次，如果有一顶点出现的次数和颜色不同的边出现的次数一样多，那么我们就能找到这一个顶点达到题目要求

#include <iostream>
#include <cstdio>
#include <algorithm>
#include <vector>
#include <cstring>
#include <queue>
#include <map>
using namespace std;
const int maxn = 2e5+100;
const int inf = 0x7fffffff;
int u[maxn],v[maxn];
int col[maxn];
int du[maxn];
int main(void)
{
    int n;
    scanf("%d",&n);
    memset(du,0,sizeof(du));
    for(int i=0;i<n-1;i++)
        scanf("%d %d",&u[i],&v[i]);
    for(int i=1;i<=n;i++)
        scanf("%d",&col[i]);
    int cnt = 0;
    for(int i=0;i<n-1;i++)
    {
        if(col[u[i]]!=col[v[i]])
        {
            cnt++;
            du[u[i]]++;
            du[v[i]]++;
        }
    }
    int flag = 1;
    for(int i=1;i<=n;i++)
    {
        if(du[i]==cnt)
        {
            flag = 0;
            puts("YES");
            printf("%d\n",i);
            break;
        }
    }
    if(flag)
        puts("NO");
    return 0;
}