C. Timofey and a tree 观察题 ＋ ｄｆｓ模拟

http://codeforces.com/contest/764/problem/C

题意：在n个顶点中随便删除一个，然后分成若干个连通子图，要求这若干个连通子图的颜色都只有一种。

记得边是双向的，wa15的可能是不知道边是双向的吧。

一个观察：如果某条边连接的两个顶点的颜色不同，那么可以看看删除这两个顶点，成立就成立，不成立就不成立。

因为必定要把这两个顶点分开。

然后就是暴力dfs了。

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cmath>
#include <algorithm>
#include <assert.h>
#define IOS ios::sync_with_stdio(false)
using namespace std;
#define inf (0x3f3f3f3f)
typedef long long int LL;


#include <iostream>
#include <sstream>
#include <vector>
#include <set>
#include <map>
#include <queue>
#include <string>
#include <bitset>

const int maxn = 2e6 + 20;
struct node {
    int u, v, tonext;
}e[maxn];
int num;
int first[maxn];
void add(int u, int v) {
    num++;
    e[num].u = u;
    e[num].v = v;
    e[num].tonext = first[u];
    first[u] = num;
}
int c[maxn], in[maxn];
int u[maxn], v[maxn];
set<int>ss;
int vis[maxn];
int DFN;
int dfs(int cur, int no) {
    ss.insert(c[cur]);
    if (ss.size() >= 2) return false;
    bool flag = true;
    for (int i = first[cur]; i; i = e[i].tonext) {
        if (e[i].v == no) continue;
        if (vis[e[i].v] == DFN) continue;
        vis[e[i].v] = DFN;
        flag = flag && dfs(e[i].v, no);
    }
    return flag;
}
bool del(int cur) {
    for (int i = first[cur]; i; i = e[i].tonext) {
        ss.clear();
        DFN++;
        vis[cur] = DFN;
        vis[e[i].v] = DFN;
        if (dfs(e[i].v, inf) == false) return false;
    }
    return true;
}
void work() {
    int n;
    cin >> n;
    for (int i = 1; i <= n - 1; ++i) {
        cin >> u[i] >> v[i];
    }
    for (int i = 1; i <= n; ++i) {
        cin >> c[i];
    }
    int which = inf;
    for (int i = 1; i <= n - 1; ++i) {
        if (c[u[i]] != c[v[i]] && which == inf) which = i;
        add(u[i], v[i]);
        add(v[i], u[i]);
    }
    if (which == inf) {
        cout << "YES" << endl;
        cout << 1 << endl;
        return;
    }
//    cout << which << "  " << root << endl;
    if (del(u[which])) {
        cout << "YES" << endl;
        cout << u[which] << endl;
        return;
    }
    if (del(v[which])) {
        cout << "YES" << endl;
        cout << v[which] << endl;
        return;
    }
    cout << "NO" << endl;
}

int main() {
#ifdef local
    freopen("data.txt", "r", stdin);
//    freopen("data.txt", "w", stdout);
#endif
    work();
    return 0;
}