【模板】树的直径

洛谷P5536 【XR-3】核心城市

dfs记录直径

#include<bits/stdc++.h>
using namespace std;
const int maxn=1e6;
int n,m,k,p,to[maxn],head[maxn],next[maxn],gen,gg,f[maxn],mid,mdeep[maxn],deep[maxn],he[maxn];
void add(int x,int y){
	next[++m]=head[x];
	to[m]=y;
	head[x]=m;
}
void dfs1(int x,int fa){
	if(deep[x]>gen){
		gen=deep[x];
		gg=x;
	}
	for(int i=head[x];i;i=next[i]){
		int y=to[i];
		if(y==fa)continue;
		deep[y]=deep[x]+1;
		dfs1(y,x);
	}
}
void dfs2(int x,int fa){
	if(deep[x]>gen){
		gen=deep[x];
		gg=x;
	}
	for(int i=head[x];i;i=next[i]){
		int y=to[i];
		if(y==fa)continue;
		f[y]=x;//记录直径
		deep[y]=deep[x]+1;
		dfs2(y,x);
	}
}
void dfs3(int x,int fa){
	mdeep[x]=deep[x];
	for(int i=head[x];i;i=next[i]){
		int y=to[i];
		if(y==fa)continue;
		deep[y]=deep[x]+1;
		dfs3(y,x);
		mdeep[x]=max(mdeep[y],mdeep[x]);
	}
}
int main(){
	int x,y;
	scanf("%d%d",&n,&k);
	for(int i=1;i<=n-1;i++){
		scanf("%d%d",&x,&y);
		add(x,y),add(y,x);
	}
	dfs1(1,0);
	memset(deep,0,sizeof(deep));
	gen=0;
	dfs2(gg,0);
	mid=gg;
	for(int i=1;i<=(gen+1)/2;i++)
	mid=f[mid];//直径中点
	memset(deep,0,sizeof(deep));
	dfs3(mid,0);
	for(int i=1;i<=n;i++)
	he[i]=mdeep[i]-deep[i];
	sort(he+1,he+1+n);
	printf("%d\n",he[n-k]+1);
}

/* 求树的直径 */
void tree_dp(int u, int fa){
for (int i = head[u]; i != -1; i = E[i].next){
int to = E[i].to;
if (fa == to)
continue;
tree_dp(to, u);
ans = max(ans, dp[u] + dp[to] + E[i].val);
dp[u] = max(dp[u], dp[to] + E[i].val);
}}