Codeforces 673 A. Timofey and a tree

Problem Statement:

http://codeforces.com/problemset/problem/763/A

Approach :
The easiest way to solve this problem is like this. First find all the edges which have nodes having dissimilar values. Let’s call these edges “Magic Edges”. These edges are basically those edges which shall connect two different sub-trees. Now, simply check whether all the “Magic Edges” have some common element. This common element shall be the required answer. If there is no such special element then print “NO”.
A special case is when all the nodes have same color. In this case the answer would be just any node.

Solution :

#include<bits/stdc++.h>
using namespace std ;
//-------------------------------
typedef long long ll ;
typedef vector<int> vi ;
typedef pair<int,int> ii ;
//-------------------------------
#define _CRT_SECURE_NO_WARNINGS
#define pb(a) push_back(a)
#define mp(a,b) make_pair(a,b)
#define rep(i,a,b) for(ll i=(a) ; (i)<(b) ; ++i)
#define inf 2000000000
#define endl "\n"
#define de(x) cerr<<#x<<" is "<<x<<endl;
#define max(a,b) ( (a>b) ? (a) : (b)  )
#define min(a,b) ( (a<b) ? (a) : (b)  )
//------------------------------
int ri() { char c = getchar(); while(c<'0' || c>'9') c=getchar(); int ret = 0; while(c>='0' && c<='9') { ret=10*ret+c-48; c=getchar(); } return ret; }

int n,u,v;
int a[211111];
vector< ii > ve,ne ;
int main()
{
    cin >> n ;
    rep(i,1,n)
    {
        cin >> u >> v ;
        ve.pb(mp(u,v)) ;
    }
    rep(i,1,n+1)
    cin >> a[i] ;
    rep(i,0,n-1)
    {
        if(a[ ve[i].first ] != a[ ve[i].second ])
        {
            ne.pb(ve[i]);
        }
    }
    if(ne.size() == 0)
    {
        cout << "YES\n1";
        return 0;

    }
    map<int,int>mo;
    rep(i,0,ne.size())
    {
        mo[ ne[i].first ]++;
        mo[ ne[i].second ]++;
    }
    for(auto x:mo)
    {
        if(mo[x.first] == ne.size())
        {
            cout << "YES\n" ;
            cout << x.first;
            return 0;
        }
    }
    cout << "NO" ;
}