55分求助
查看原帖
55分求助
497368
0616allen楼主2023/2/14 16:20

dfs是搜索求直径,query_dm是树上dp求直径,dfs2是把路径上的边权改为-1。

#include<bits/stdc++.h>
using namespace std;
const int N=1e5+5;
int n,k,p,fa[N],maxi,dis[N];
struct Node{
	int nxt,w;
};
vector<Node> ve[N];
void query_dm(int cur,int fa)
{
	for(int i=0;i<ve[cur].size();i++){
		int next=ve[cur][i].nxt;
		int w=ve[cur][i].w;
		if(next==fa){
			continue;
		}
		query_dm(next,cur);
		maxi=max(maxi,dis[cur]+dis[next]+w);
		dis[cur]=max(dis[cur],dis[next]+w);
	}
	return ;
}
void dfs(int cur,int ffa,int dis,int num)
{
	if(num==2){
		fa[cur]=ffa;
	}
	for(int i=0;i<ve[cur].size();i++){
		int next=ve[cur][i].nxt;
		int w=ve[cur][i].w;
		if(next==ffa){
			continue;
		}
		dfs(next,cur,dis+w,num);
	}
	if(maxi<dis){
		p=cur;
		maxi=dis;
	}
	return ;
}
void dfs2(int cur)
{
	if(fa[cur]==0){
		return ;
	}
	for(int i=0;i<ve[cur].size();i++){
		int next=ve[cur][i].nxt;
		if(next==fa[cur]){
			ve[cur][i].w=-1;
			dfs2(next);
		}
	}
	return ;
}
int main()
{
	cin>>n>>k;
	for(int i=1;i<n;i++){
		int x,y;
		cin>>x>>y;
		ve[x].push_back({y,1});
		ve[y].push_back({x,1});
	}
	if(k==1){
		query_dm(1,0);
		cout<<2*n-maxi-1;
	}
	else{
		dfs(1,0,0,1);
		maxi=0;
		dfs(p,0,0,2);
		int tmp=maxi;
		dfs2(p); 
		maxi=0;
		query_dm(1,0);
		cout<<2*n-maxi-tmp;
	}
}
2023/2/14 16:20
加载中...