codeforces 764 C. Timofey and a tree（dfs+思维）

题目链接：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;
}