Codeforces 764C - Timofey and a tree（dfs）

题意

给一棵树，树上每个节点都有颜色，现在要求删掉一个点，在删掉这个点后，剩下的所有子树，每个子树上的所有点颜色相同。

思路

缩点。
我们将树上相连的，并且颜色相同的点缩成一个点。在缩点重新建图后，假设重新建图后有tot个点，那么一定要有一个点满足度为tot - 1（即这个点和其他所有点相连并且其他点互不相连）。

代码

#include <bits/stdc++.h>

using namespace std;

inline int in() {int x; scanf("%d", &x); return x;}
#define pr(x) {cout << #x << ' ' << x << endl;}

const int maxn = 100000 + 5;
vector<int> G[maxn];
vector<int> R[maxn];
int c[maxn], cc[maxn], n;
int vis[maxn], mc = 0, tot = 0, code[maxn];

void dfs(int u, int p) {
    for (auto v : G[u]) {
        if (v == p) continue;
        if (c[u] != c[v]) {
            if (!code[u]) {code[u] = ++tot; vis[tot] = u;}
            if (!code[v]) {code[v] = ++tot; vis[tot] = v;}
            R[code[u]].push_back(code[v]);
            R[code[v]].push_back(code[u]);
        }
        dfs(v, u);
    }
}


int main() {
    n = in();
    for (int i = 0; i < n - 1; i++) {
        int u = in(), v = in();
        G[u].push_back(v);
        G[v].push_back(u);
    }
    for (int i = 1; i <= n; i++) {
        c[i] = in();
    }
    dfs(1, -1);
    if (tot == 0) {
        cout << "YES" << endl;
        cout << 1 << endl;
        return 0;
    }
    for (int i = 1; i <= tot; i++) {
        if (R[i].size() == tot - 1) {
            cout << "YES" << endl;
            cout << vis[i] << endl;
            return 0;
        }
    }
    cout << "NO" << endl;
    return 0;
}