Codeforces Round #395 (Div. 2) - C

 

题目链接：http://codeforces.com/contest/764/problem/C

题意：给定一个n个点n-1条边的树，然后树上每一个点都有一个颜色。现在问你能不能删除一个点是个删除这个点之后所新产生的多棵中每棵树中点的颜色都一样。如果存在合法的删除点输出YES和该点。否则输出NO。

思路：起初我暴力枚举删除n个点然后判断发现tle了。 考虑下题目要求删除后新产生的多棵树中每个点的颜色要一致，那么我们只要枚举输入的边发现边对应的两个点的颜色不一样那么肯定要删除其中一个点，因为假设点x和点y颜色不一样并且有一边x-y的边，那么不删除x/y的话，新产生的多棵树中x和y一定在同一棵树上，因为x-y有边而且你只能删除一个点并且的删除的点不是x和y，那么x和y在同个树上。所以只需枚举两个点即可。 注意有一种情况就是输入的所有边的两个点颜色都一样，这种情况随便删除哪个点都是合法的。

#define _CRT_SECURE_NO_DEPRECATE
#include<iostream>
#include<cstring>
#include<string>
#include<algorithm>
#include<stdio.h>
#include<queue>
#include<vector>
#include<stack>
#include<map>
#include<set>
#include<time.h>
#include<cmath>
using namespace std;
#define x first
#define y second
#define pb push_back
#define mp make_pair
typedef long long int LL;
const int inf = 0x3f3f3f3f;
const LL INF = 0x3f3f3f3f3f3f3f3fLL;
const int MAXN = 1e5 + 10;
int c[MAXN];
bool flag;
vector<int>G[MAXN];
void dfs(int u, int fa, int col){
    if (!flag){ return; }
    for (int i = 0; i < G[u].size(); i++){
        if (G[u][i] != fa){
            if (c[G[u][i]] != col){
                flag = false;
            }
            else{
                dfs(G[u][i], u, col);
            }
        }
    }
}
int solve(int x){
    flag = true;
    for (int i = 0; i < G[x].size()&&flag; i++){
        dfs(G[x][i], x, c[G[x][i]]);
    }
    if (flag){
        return x;
    }
    else{
        return -1;
    }
}
int main(){
//#ifdef kirito
//    freopen("in.txt", "r", stdin);
//    freopen("out.txt", "w", stdout);
//#endif
//    int start = clock();
    int n;
    while (~scanf("%d", &n)){
        for (int i = 0; i <= n; i++){
            G[i].clear();
        }
        for (int i = 1; i < n; i++){
            int u, v;
            scanf("%d%d", &u, &v);
            G[u].push_back(v);
            G[v].push_back(u);
        }
        for (int i = 1; i <= n; i++){
            scanf("%d", &c[i]);
        }
        int ans = -1; bool k = false;
        for (int i = 1; i <= n&&k==false; i++){
            for (int j = 0; j < G[i].size(); j++){
                if (c[i] != c[G[i][j]]){ //有边并且颜色不一样
                    ans = solve(i);
                    if (ans == -1){
                        ans = solve(G[i][j]);
                    }
                    k = true;
                    break;
                }
            }
        }
        if (k == false){
            ans = 1;
        }
        if (ans != -1){
            printf("YES\n%d\n", ans);
        }
        else{
            printf("NO\n");
        }
    }
//#ifdef LOCAL_TIME
//    cout << "[Finished in " << clock() - start << " ms]" << endl;
//#endif
    return 0;
}