考察并查集以及树的搜索
#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;
}