CF - 764C. Timofey and a tree - 记忆化dfs判重/并查集

1.题目描述：

C. Timofey and a tree

time limit per test

2 seconds

memory limit per test

256 megabytes

input

standard input

output

standard output

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

input

3
1 2
2 3
1 2 3

output

YES
2

input

4
1 2
2 3
3 4
1 2 1 2

output

NO

2.题意概述：

给你一棵树，你可以以任意一点为根节点，但是要求他的所有子树满足整棵子树颜色相同，问你能否找到这样的根节点，能则输出任意一个，不能则NO。

3.解题思路：

最开始纯粹暴力dfs，T在第九个，然后考虑记忆化搜索，如果当前点不满足，那么今后搜索到这个点肯定也是不满足的，就没有必要搜索下去了。惊天的发现，C++14照样T在第七个测试点，但是换GUN C++居然过了，看来以后还是少用Clang。

还有一种做法是并查集：

先建图，遍历每个结点，统计每个结点的度，对结点i，若其相邻的结点j与其颜色相同，那么将其合并，并将结点i和结点j的度都减1，合并完后，所有颜色相同且相互邻接的点合并为一个集合，统计集合个数cnt，将一个集合看作一个点，在这个虚图中满足条件的点i一定满足cnt-1==deg[i]（树的性质），在合并后的原图中所有点中寻找满足cnt-1==deg[i]的点。

4.AC代码：

dfs:

#include <bits/stdc++.h>
#define INF 0x3f3f3f3f
#define maxn 100010
#define lson root << 1
#define rson root << 1 | 1
#define lent (t[root].r - t[root].l + 1)
#define lenl (t[lson].r - t[lson].l + 1)
#define lenr (t[rson].r - t[rson].l + 1)
#define N 1111
#define eps 1e-6
#define pi acos(-1.0)
#define e exp(1.0)
using namespace std;
const int mod = 1e9 + 7;
typedef long long ll;
typedef unsigned long long ull;
vector<int> g[maxn];
int n, col[maxn];
bool flag, vis[maxn];
map<int, bool> can[maxn];
inline bool scan_d(int &num)
{
	char in; bool IsN = false;
	in = getchar();
	if (in == EOF) return false;
	while (in != '-' && (in<'0' || in>'9')) in = getchar();
	if (in == '-') { IsN = true; num = 0; }
	else num = in - '0';
	while (in = getchar(), in >= '0'&&in <= '9') {
		num *= 10, num += in - '0';
	}
	if (IsN) num = -num;
	return true;
}
void dfs(int u, int color)
{
	int sz = g[u].size();
	for (int i = 0; i < sz && flag; i++)
	{
		int v = g[u][i];
		if (!vis[v])
		{
			if (can[u][v] && color)
			{
				flag = 0;
				return;
			}
			if (color == 0 || (color != 0 && col[v] == color))
			{
				vis[v] = 1;
				dfs(v, col[v]);
				if (!flag)
				{
					can[u][v] = 1;
					return;
				}
				vis[v] = 0;
			}
			else
			{
				can[u][v] = 1;
				flag = 0;
				return;
			}
		}
	}
}
int main()
{
#ifndef ONLINE_JUDGE
	freopen("in.txt", "r", stdin);
	freopen("out.txt", "w", stdout);
	long _begin_time = clock();
#endif
	while (~scanf("%d", &n))
	{
		for (int i = 1; i <= n; i++)
		{
			g[i].clear();
			can[i].clear();
		}
		for (int i = 0; i < n - 1; i++)
		{
			int u, v;
			scanf("%d%d", &u, &v);
			g[u].push_back(v);
			g[v].push_back(u);
		}
		for (int i = 1; i <= n; i++)
			scanf("%d", &col[i]);
		int ans = 0;
		for (int i = 1; i <= n; i++)
		{
			memset(vis, 0, sizeof(vis));
			vis[i] = 1;
			flag = 1;
			dfs(i, 0);
			if (flag)
			{
				ans = i;
				break;
			}
		}
		if (!ans)
			puts("NO");
		else
			printf("YES\n%d\n", ans);
	}
#ifndef ONLINE_JUDGE
	long _end_time = clock();
	printf("time = %ld ms.", _end_time - _begin_time);
#endif
	return 0;
}

并查集：

#include <bits/stdc++.h>
#define INF 0x3f3f3f3f
#define maxn 100100
#define lson root << 1
#define rson root << 1 | 1
#define lent (t[root].r - t[root].l + 1)
#define lenl (t[lson].r - t[lson].l + 1)
#define lenr (t[rson].r - t[rson].l + 1)
#define N 1111
#define eps 1e-6
#define pi acos(-1.0)
#define e exp(1.0)
using namespace std;
const int mod = 1e9 + 7;
typedef long long ll;
typedef unsigned long long ull;
int pa[maxn], r[maxn], col[maxn], deg[maxn];
vector<int> g[maxn];
inline void init(int n)
{
	for (int i = 1; i <= n; i++)
	{
		pa[i] = i;
		r[i] = 1;
		deg[i] = 0;
		g[i].clear();
	}
}
int findset(int x)
{
	if (pa[x] == x)
		return x;
	return pa[x] = findset(pa[x]);
}
inline void unionset(int a, int b)
{
	int a1 = findset(a);
	int b1 = findset(b);
	if (a1 != b1)
	{
		pa[a1] = b1;
		r[b1] += r[a1];
	}
}
int main()
{
#ifndef ONLINE_JUDGE
	freopen("in.txt", "r", stdin);
	freopen("out.txt", "w", stdout);
	long _begin_time = clock();
#endif
	int n;
	while (~scanf("%d", &n))
	{
		init(n);
		for (int i = 0; i < n - 1; i++)
		{
			int u, v;
			scanf("%d%d", &u, &v);
			deg[u]++;
			deg[v]++;
			g[u].push_back(v);
			g[v].push_back(u);
		}
		for (int i = 1; i <= n; i++)
			scanf("%d", &col[i]);
		for (int i = 1; i <= n; i++)
		{
			int sz = g[i].size();
			for (int j = 0; j < sz; j++)
				if (col[i] == col[g[i][j]])
				{
					if (findset(i) != findset(g[i][j]))
					{
						unionset(i, g[i][j]);
						deg[i]--;
						deg[g[i][j]]--;
					}
				}
		}
		int cnt = 0;
		for (int i = 1; i <= n; i++)
			if (pa[i] == i)
				cnt++;
		int ans = 0;
		for (int i = 1; i <= n; i++)
			if (deg[i] == cnt - 1)
			{
				ans = i;
				break;
			}
		if (ans)
			printf("YES\n%d\n", ans);
		else
			puts("NO");
	}
#ifndef ONLINE_JUDGE
	long _end_time = clock();
	printf("time = %ld ms.", _end_time - _begin_time);
#endif
	return 0;
}