[Codeforces] 764C.(763A.) Timofey and a tree

題目連結:http://codeforces.com/contest/764/problem/C

---

每次都找degree為1的人，看他跟他連接的那個人顏色一不一樣，如果依樣就把那條邊刪掉，直到每個degree為一的都找過為止，然後看是不是只有一個人degree不為一或零，如果有，他就是唯一的答案。但是要注意一下，因為有可能沒有唯一解，因此需要特判最後若degree度和為零(所有人都被砍掉了)，那一定是因為所有人顏色一樣，所以隨便輸出一個範圍內的數就好了；degree和為二的時候，表示最後只剩兩點，那答案也必定是那兩點之一，隨便輸出一個就好了。

#include <bits/stdc++.h>
using namespace std;
#define N 100000

set<int> tree[N+5];

int color[N + 5];
int deg[N + 5]={0};

int main(){
	int n;
	scanf("%d",&n);
	for(int i=0;i<n-1;i++){
		int a,b;
		scanf("%d%d",&a,&b);
		tree[a].insert(b);
		tree[b].insert(a);
		deg[a]++;
		deg[b]++;
	}
	for(int i=1;i<=n;i++)
		scanf("%d",&color[i]);
	queue<int> deg1;
	for(int i=1;i<=n;i++)
		if(deg[i]==1)deg1.push(i);
	while(!deg1.empty()){
		int top=deg1.front();
		deg1.pop();
		if(deg[top]!=1)continue;
		int f = *tree[top].begin();
		if(color[f]==color[top]){
			tree[f].erase(top);
			deg[top]--;
			deg[f]--;
			if(deg[f]==1) deg1.push(f);
		}
	}
	int cnt=0,sum=0;
	int ans;
	for(int i=1;i<=n;i++){
		sum+=deg[i];
		if(deg[i]>1){
			cnt++;
			ans=i;
		}
	}
	if(cnt<=1){
		puts("YES");
		if(sum == 0){
			printf("1");
		}else if(sum==2){
			for(int i=1;i<=n;i++){
				if(deg[i]==1){
					printf("%d",i);
					break;
				}
			}
		}else{
			printf("%d",ans);
		}
	}else{
		puts("NO");
	}
	return 0;
}