CF763A Timofey and a tree（思维）

现在有一个无根树，我们加入把其中的一个结点看做根，那么这个根上的每一个子树中都不能有一个结点的颜色不同，如果都相同，那么此时就是合法的，我们输出YES及这个结点。如果不存这样的根，那么输出NO。

思路：
因为如果一个树的子树中只要有一个点与这个子树中的其他点颜色不同，那么这样主角就会生气，所以这样只有一种情况是合法的，就是当这个无根树的某个结点作为根时，其连接到的结点与根的颜色不一样的个数等于所有结点与这个根颜色不同的结点个数。(有点绕，耐心理解^-)。因为子树中一旦出现某个结点与其他结点颜色不同就会造成整个树不合法。

#include <bits/stdc++.h>
using namespace std;
const int maxn = 1e5 + 5;
int a[maxn], b[maxn], cnt[maxn], col[maxn];
int main() {
	int n;
	cin >> n;
	for (int i = 1; i <= n - 1; i++) {
		cin >> a[i] >> b[i];
	}
	for (int i = 1; i <= n; i++) {
		cin >> col[i];
	}
	int res = 0;
	for (int i = 1; i <= n - 1; i++) {
		if (col[a[i]] != col[b[i]]) { //相邻节点的颜色不同（）
			cnt[a[i]]++; cnt[b[i]]++; res++;
		}
	}
	for (int i = 1; i <= n; i++) {
		if (cnt[i] == res) {
			cout << "YES" << endl;
			cout << i << endl;
			return 0;
		}
	}
	cout << "NO" << endl;
	return 0;
}