Comet OJ - Contest #9 & X Round 3 【XR-3】核心城市(树的直径）

题目链接：https://cometoj.com/contest/64/problem/C?problem_id=3520

思路：比赛的时候想到树的直径，但是不知道如何选点。。。。选树的中心：树的直径的中点为起始点，优先加入子树最远距离大的点，最后再统计一遍选过k个点以后每个子树的最远距离即可。
代码有点low，用了4次dfs。。

代码：

#include <bits/stdc++.h>

using namespace std;

const int maxn = 1e5 + 5;

vector <int> g[maxn];
int d[maxn] , mnext[maxn]; 
bool vis[maxn];
vector <int> vex;
int mm , p;

struct node {
	int v;
	node(int a):v(a){}
	bool operator < (const node b)const {
		return d[v] < d[b.v];
	}
};

void dfs(int u , int dist) {
	vis[u] = 1;
	if(dist > mm) {
		p = u;
		mm = dist;
	}
	for(int i = 0 ; i < g[u].size() ; i++) {
		int v = g[u][i];
		if(!vis[v]) {
			dfs(v , dist + 1);
			if(d[u] < d[v] + 1) {
				d[u] = d[v] + 1;
				mnext[u] = v;		
			} 
		}
	}
	
}

int main() {
	int n , k;
	ios::sync_with_stdio(0);
	cin >> n >> k;
	for(int i = 0 ; i < n - 1; i++) {
		int u , v;
		cin >> u >> v;
		g[u].push_back(v);
		g[v].push_back(u);
	}
	mm = 0; 
	dfs(1 , 0);
	int p1 = p;
	memset(vis,0,sizeof(vis));memset(d,0,sizeof(d));memset(mnext,0,sizeof(mnext));
	dfs(p1 , 0);
	int t = mnext[p1];
	vex.push_back(p1);
	while(t) {
		vex.push_back(t);
		t = mnext[t];
	}
	int u = vex[vex.size() / 2 ];
	memset(vis,0,sizeof(vis));memset(d,0,sizeof(d));
	dfs(u , 0);
	priority_queue <node> q;
	if(k == 1)cout << d[u] << " ";
	else {
		memset(vis,0,sizeof(vis));
		int cnt = 1;
		q.push(node(u));
		while(!q.empty()) {
			node t = q.top(); q.pop();
			if(vis[t.v])continue;
			vis[t.v] = 1;
			if(cnt >= k)break;
			cnt++;
			for(int i = 0 ; i < g[t.v].size() ; i++) {
				int v = g[t.v][i];
				if(!vis[v]) {
					q.push(node(v));
				}
			}
		}
		mm = 0;
		for(int i = 1 ; i <= n ; i++) {
			if(vis[i])dfs(i , 0);
		}
		cout << mm << "\n";
	} 
	return 0;
}