Codeforces Round #395 (Div. 2) C. Timofey and a tree

Each New Year Timofey and his friends cut down a tree of n vertices and bring it home. After that they paint all the n its vertices, so that the i-th vertex gets color ci.

Now it's time for Timofey birthday, and his mother asked him to remove the tree. Timofey removes the tree in the following way: he takes some vertex in hands, while all the other vertices move down so that the tree becomes rooted at the chosen vertex. After that Timofey brings the tree to a trash can.

Timofey doesn't like it when many colors are mixing together. A subtree annoys him if there are vertices of different color in it. Timofey wants to find a vertex which he should take in hands so that there are no subtrees that annoy him. He doesn't consider the whole tree as a subtree since he can't see the color of the root vertex.

A subtree of some vertex is a subgraph containing that vertex and all its descendants.

Your task is to determine if there is a vertex, taking which in hands Timofey wouldn't be annoyed.

Input

The first line contains single integer n (2 ≤ n ≤ 105) — the number of vertices in the tree.

Each of the next n - 1 lines contains two integers u and v (1 ≤ u, v ≤ n, u ≠ v), denoting there is an edge between vertices u and v. It is guaranteed that the given graph is a tree.

The next line contains n integers c1, c2, ..., cn (1 ≤ ci ≤ 105), denoting the colors of the vertices.

Output

Print "NO" in a single line, if Timofey can't take the tree in such a way that it doesn't annoy him.

Otherwise print "YES" in the first line. In the second line print the index of the vertex which Timofey should take in hands. If there are multiple answers, print any of them.

Examples

input

4
1 2
2 3
3 4
1 2 1 1

output

YES
2

首先庆祝第一次做了3题，而且这个第三题在我看来还不是很水，还是挺考验思维的。

开始我读了几遍题，都没搞懂题意，再三读题，原来就是：

“Your task is to determine if there is a vertex, taking which in hands Timofey wouldn't be annoyed.”这句话。

题意：给定一棵树，每个节点有一种颜色，问是否存在一个点，去掉后，得到的子树的颜色都是纯正的，即

每棵子树上都只有一种颜色，树之间的颜色可以不同。存在的话打印出YES跟对应节点，答案可以不同。

否则打印NO。

思路历程：这个题意的直接思路就是枚举要删除哪个点，对于每个点，检查拆出来的子树颜色是否只有一种。

统计子树颜色种类就是要dfs，每个都这么做20万个点会超时，就想到了树形动态规划里面的二次dfs，- -，郁闷了半小

时都没想出可行方法。该方法作罢。

直接的方法不行那就换个方法，

思路：改成添加哪个点可以把所有同颜色的都连接成一棵完整的树。那么这个点被移除时，

也就可以得到颜色纯正的子树了。

首先dfs用并查集，把树上直接相连且某种颜色一致的节点加到一个集合里面，这样就得到了不同纯色子树构成的集合。构建并查集时顺便统计每个集合的个数。

然后枚举每个节点，看通过它连接得到树的总节点数是否等于原始树节点数n，如果是那就找到了要求的点。

复杂度分析：dfs复杂度O(n)，枚举节点时要遍历与当前节点直连的节点，其实就是遍历边，一条边有连两个节点，

所以边最多被遍历两遍，边数是n-1，所以复杂度还是O(n)。总复杂度就是O(n)。

#include <bits/stdc++.h>
using namespace std;
const int N = 200010;
vector<int> rec[N];
int col[N];
int fa[N];
int cnt[N];
bool vis[N];
void init(int n)
{
	for (int i = 1; i <= n; i++)
		fa[i] = i, cnt[i] = 1;
}
int find(int x)
{
	return fa[x] == x?x:(fa[x] = find(fa[x]));
}
void merge(int u, int v)
{
	u = find(u);
	v = find(v);
	if (u != v)
	{
		fa[u] = v;
		cnt[v] += cnt[u];
	}	
}
void dfs(int u, int pre)
{
	if (vis[u])
		return;
	vis[u] = true;
	for (int i = 0; i < rec[u].size(); i++)
	{
		int v = rec[u][i];
		if (v == pre)
			continue;
		
		if (col[v] == col[u])
			merge(v, u);
		dfs(v, u);
	}
}
int main() 
{
  int n;
  while (~scanf("%d", &n))
  {
  	for (int i = 1; i < n; i++)
  	{
  		int a, b;
  		scanf("%d%d", &a, &b);
  		rec[a].push_back(b);
  		rec[b].push_back(a);
  	}
  	for (int i = 1;i <= n; i++)
  		scanf("%d", col + i);
  	init(n);
  	memset(vis, 0, sizeof(vis));
  	dfs(1, -1);

  	int ans = -1;
  	for (int i = 1; i <= n; i++)
  	{
  		int tmp = cnt[find(i)];
  		for (int j = 0; j < rec[i].size(); j++)
  		{
  			int v = rec[i][j];
  			if (find(v) != find(i))
  				tmp += cnt[find(v)];
  		}
  		if (tmp == n)
  		{
  			ans = i;
  			break;
  		}
  	}
  	if (~ans)
  	{
  		puts("YES");
  		printf("%d\n", ans);
  	}else
  		puts("NO");
  }
  return 0;
}

看了一下别人的代码，发现其实有更简单的思路：

如果选中一个点去掉，那么剩下子树里面每条边连接的点的颜色应该是一样的这样子树颜色才能纯正。

那就统计一下两个点颜色不同的边的数目，同时统计边里面包含的点的次数。

如果一个点出现的次数与颜色不同的边的次数一样，说明这个点是这些边的公共点，去掉就可以得到颜色纯正的子树

了。

n = input()
a = {}
vs = []

edge = 0
for i in xrange(1, n):
	a[i] = 0
	u, v = map(int, raw_input().split(' '))
	vs.append((u,v))
	
a[n] = 0
col = map(int, raw_input().split(' '))
for x in vs:
	if col[x[0] - 1] != col[x[1] - 1]:
		edge += 1
		a[x[0]] += 1
		a[x[1]] += 1
for i in xrange(1, n + 1):
	if a[i] == edge:
		print 'YES'
		print i
		exit(0)
print 'NO'