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

C. Timofey and a tree

链接：

http://codeforces.com/contest/764/problem/C

题意: 

一棵树中各个节点被染上了c[i]颜色; 
让你在一棵树中随便选一个节点作为根节点,然后把整棵树抬起来; 
问你是否存在一个根节点，这个根节点的直系儿子节点的子树里面的所有节点的颜色都一样; 

题解：

要理解树的构造以及特点，每棵树的子树是不相连的，就是几个分块了，所以这个点要想使所有子树都是一种颜色，其实就是这个点所连边包括所有端点不同色的边，因为还有端点不同色的边没有与根节点相连，那么肯定某一颗子树含有这个边，也就是这颗子树并没有同色，因为子树之间没有边相连，所以只能在某一科子树里。所以做法就是记录所有不同色的边的个数，记录每个点连接的不同色的边的个数，如果某个两者相等说明这个点连着所有不同色的边，也就是符合题意了

代码：

#include <iostream>  
#include <cstdio>  
#include <algorithm>  
using namespace std;
const int maxn = 2e5 + 5;
int u[maxn], v[maxn], cnt[maxn], c[maxn], sum = 0;
int main()
{
    int n;
    cin >> n;
    for (int i = 1; i < n; i++)
        cin >> u[i] >> v[i];
    for (int i = 1; i <= n; i++)
        cin >> c[i];
    for (int i = 1; i < n; i++)
    {
        if (c[u[i]] != c[v[i]])
            sum++, cnt[u[i]]++, cnt[v[i]]++;

    }
    for (int i = 1; i <= n; i++)
    {
        if (cnt[i] == sum)
        {
            cout << "YES\n" << i << endl;
            return 0;
        }
    }
    cout << "NO" << endl;
    return 0;
}