【CodeForces - 764C】Timofey and a tree（思維）

 題意：找一個根節點使得他的每顆子樹的所有結點顏色一樣。
 如果有兩個結點的顏色不一樣，並且他們之間還有邊的話，那麼根節點一定在他倆之間，所以對這兩個結點進行判斷即可

ac代碼：

#include<iostream>
#include<cstdio>
#include<algorithm>
#include<cmath>
#include<cstring>
#include<queue>
#include<stack>
#include<map>
#include<set>
#include<string>
#include<vector>
#include<unordered_map>
#define mod (1000000007)
using namespace std;
typedef long long ll;
const int maxn = 2e5 + 5;
const int inf = 0x3f3f3f3f;
vector<int> vv[maxn];
int u,v,c[maxn];
int dfs(int x,int fa){
    for(auto ty:vv[x]){
        if(ty==fa) continue;
        int tt=dfs(ty,x);
        if(tt==0) return 0;
        if(fa==-1){
            ;
        }
        else{
            if(c[x]!=c[ty]) return 0;
        }
    }
    return 1;
}
int main() {
    int n,r1=-1,r2=-1;
    scanf("%d",&n);
    for(int i=1;i<n;i++){
        scanf("%d%d",&u,&v);
        vv[v].push_back(u);
        vv[u].push_back(v);
    }
    for(int i=1;i<=n;i++) scanf("%d",&c[i]);
   for(int i=1;i<=n;i++){
       for(auto v:vv[i]){
           if(c[v]!=c[i]){
               r1=v;r2=i;break;
           }
       }
       if(r1!=-1) break;
   }
   if(r1==-1&&r2==-1)printf("YES\n1\n");
   else if(r1!=-1&&dfs(r1,-1)) printf("YES\n%d\n",r1);
   else if(r2!=-1&&dfs(r2,-1)) printf("YES\n%d\n",r2);
   else puts("NO");
    return 0 ;
}