【CF763A】 A. Timofey and a tree（贪心）

蒂莫菲和一棵树

题意：

• 给定一棵树，树中任意一个点都有一个颜色。问能否找到一个点，以这个点为根节点产生的树中所有子树（不包括根节点这颗树）的结点颜色都相同。

思路：

• 可以 dp 写也可以贪心写，但 dp 维护起来很麻烦…

• 贪心的方法是观察性质，如果 以根节点产生的树中所有子树的结点颜色都相同，那各个子树内包含的都是两端点颜色相同的边，那两端点颜色不同的边，显然就在根节点与子树之间。

• 所以我们只需要统计任意一个点连出去的边，有多少条颜色不同，和给定的树中颜色不同的边数比较，如果该点满足相等，自然合法。

$Code:$

#include<bits/stdc++.h>
#include<unordered_map>
#define mem(a,b) memset(a,b,sizeof a)
#define cinios (ios::sync_with_stdio(false),cin.tie(0),cout.tie(0))
#define sca scanf
#define pri printf
#define forr(a,b,c) for(int a=b;a<=c;a++)
#define rfor(a,b,c) for(int a=b;a>=c;a--)
#define all(a) a.begin(),a.end()
#define oper(a) (operator<(const a& ee)const)
#define endl "\n"
using namespace std;
 
typedef long long ll;
typedef pair<int, int> PII;
 
double DNF = 1e17;
const int N = 100010, M = 200010, MM = 110;
int INF = 0x3f3f3f3f, mod = 1e9 + 7;
ll LNF = 0x3f3f3f3f3f3f3f3f;
int n, m, k, T, S, D, K;
int col[N], dif[N];
struct edge
{
	int a, b;
}ed[N];
 
void solve() {
	cin >> n;
	forr(i, 1, n - 1) {
		int a, b;
		cin >> a >> b;
		ed[i] = { a,b };
	}
	forr(i, 1, n)cin >> col[i];
 
	int cnt = 0;
	forr(i, 1, n - 1) {
		int a = ed[i].a, b = ed[i].b;
		if (col[a] != col[b]) {
			dif[a]++;//两端点颜色不同
			dif[b]++;//对两个点来说出去的颜色不同的边数都++
			cnt++;//颜色不同的总边数
		}
	}
	
	forr(i, 1, n) {
		if (dif[i] == cnt) { //判断
			cout << "YES\n" << i;
			return;
		}
	}
	cout << "NO";
}
 
int main() {
	cinios;
 
	T = 1;
	while (T--)solve();
 
	return 0;
}
/*
*/