【zzuliOJ】1900 - 985的“树”难题（bfs & 并查集）

#include <stdio.h>
#include <cstring>
#include <vector>
#include <algorithm>
using namespace std;
#define CLR(a,b) memset(a,b,sizeof(a))
#define INF 0x3f3f3f3f
#define LL long long
#define MAX 10000
int f[MAX+11];
int fa[MAX+11];
int ch[MAX+11];
bool flag;		//记录是否是树 
vector<int> edge[MAX+11];
void init(int n)
{
	flag = true;
	for (int i = 0 ; i <= n ; i++)
	{
		edge[i].clear();
		f[i] = i;
		ch[i] = 0;
		fa[i] = -1;
	}
}
int find(int x)
{
	if (x != f[x])
		f[x] = find (f[x]);
	return f[x];
}
bool join(int x,int y)
{
	int fx,fy;
	fx = find(x);
	fy = find(y);
	if (fx != fy)
	{
		f[fx] = fy;
		return true;
	}
	return false;
}
void bfs(int root,int father)
{
	if (edge[root].size() == 1 && edge[root][0] == father)		//如果只有父节点一条线，那么它是叶子
	{
		ch[root] = 0;
		return;		//递归终止 
	}
	for (int i = 0 ; i < edge[root].size() ; i++)		//否则搜索子节点 
	{
		int node = edge[root][i];
		if (node != father)
		{
			bfs(node,root);
			ch[root] += ch[node] + 1;		//加上子节点的孩子
			fa[node] = root;		//子节点的父亲是它 
		}
	}
}
int main()
{
	int u;
	int n,root;
	scanf ("%d",&u);
	while (u--)
	{
		scanf ("%d %d",&n,&root);
		init(n);
		for (int i = 1 ; i < n ; i++)
		{
			int t1,t2;
			scanf ("%d %d",&t1,&t2);
			edge[t1].push_back(t2);
			edge[t2].push_back(t1);
			if (!join(t1,t2))		//出现了环 
				flag = false;
		}
		if (!flag)
		{
			printf ("NO\n");
			continue;
		}
		int ant = 0;
		for (int i = 1 ; i <= n ; i++)
		{
			if (f[i] == i)
				ant++;
			if (ant > 1)
			{
				flag = false;
				break;
			}
		}
		if (!flag)
		{
			printf ("NO\n");
			continue;
		}
		//以上判断是否是树
		bfs(root,root);
		fa[root] = root;
		printf ("YES\n");
		for (int i = 1 ; i < n ; i++)
			printf ("%d ",ch[i]);
		printf ("%d\n",ch[n]);
		for (int i = 1 ; i < n ; i++)
			printf ("%d ",fa[i]);
		printf ("%d\n",fa[n]);
	}
	return 0;
}