1021. Deepest Root (25)

考察并查集以及树的搜索

#include<iostream>
#include<vector>
#include<algorithm>
#include<queue>
#include<set>
//#include<_dbdao.h>
std::vector<bool> visited;
std::vector<int> step;
std::vector<std::vector<int>> edge;
std::vector<int> p;
int n;
int BFS(int s)
{
	visited.assign(n+1, false);
	std::queue<int> q;
	q.push(s);
	visited[s] = true;
	int cnt = 1;//
	step.assign(n+1, 0);
	step[s] = 1;
	int mmax = step[s];
	while(!q.empty())
	{
		int c = q.front();
		q.pop();
		for(int i = 0; i < edge[c].size(); ++i)
		{
			int index = edge[c][i];
			if(!visited[index])
			{
				q.push(index);
				step[index] = step[c]+1;
				mmax = std::max(step[index], mmax);
				visited[index] = true;
				cnt++;
				if(cnt == n) break;
			}
		}
	}
	return mmax;
}
void InitSet()
{
	p.resize(n+1);
	for(int i = 0; i < p.size(); ++i)
		p[i] = i;
}
void CompressSet(int top, int x)
{
	if(p[x] != top)
	{
		CompressSet(top, p[x]);//go to root
		p[x] = top;//when root back, then assign to top
	}
}
int FindSet(int x)
{
	if(x != p[x])
	{
		int t = FindSet(p[x]);//go to root
		CompressSet(t, x);//when root back, then compress
	}
	return p[x];
}
void UnionSet(int x, int y)
{
	int a = FindSet(x);
	int b = FindSet(y);
	p[a] = b;
}
int main()
{
	while(scanf("%d",&n)!=EOF)
	{
		edge.clear();
		edge.resize(n+1);
		for(int i = 1; i <= n-1; ++i)
		{
			int a, b;
			scanf("%d%d",&a,&b);
			edge[a].push_back(b);
			edge[b].push_back(a);
		}
		//check if it is a tree
		//construct disjoint set
		InitSet();
		for(int i = 1; i <= n; ++i)
		{
			for(int j = 0; j < edge[i].size(); ++j)
			{
				UnionSet(i, edge[i][j]);
			}
		}
		std::set<int> root;
		for(int i = 1; i <= n; ++i)
			root.insert(FindSet(i));
		if(root.size() != 1)
			printf("Error: %d components\n",root.size());
		else
		{
			//bfs
			//get one endpoint first
			std::set<int> ans;
			int mmax = BFS(1);
			for(int i = 1; i <= n; ++i)
				if(step[i] == mmax) ans.insert(i);
			//get the diameter of the tree
			mmax = BFS(*ans.begin());
			for(int i = 1; i <= n; ++i)
				if(step[i] == mmax) ans.insert(i);
			//output
			std::set<int>::iterator it;
			for(it = ans.begin(); it != ans.end(); it++)
				printf("%d\n",*it);
		}
		
	}

	return 0;
}