洛谷 4149 [IOI2011]Race （dsu on tree）

题目链接：https://www.luogu.com.cn/problem/P4149

题意：给一棵树，n个点，每条边有权。求一条简单路径，权值和等于 k，且边的数量最小。

思路：将问题转化为子树中，对于节点u，可以把u看作简单路径两个端点v1和v2的LCA。预处理每个节点到根的距离（数组dis）和深度(数组dep)，问题转化为求最小的dep[v1]+dep[v2]-2*dep[u]使得dis[v1]+dis[v2]-2*dis[u]==k 。维护子树中每个dis的最小深度。在计算u子树的答案时，要先更新答案，加上u以及轻儿子v的贡献之前。这样是为了避免将v的子树中简单路径（v1->v2）统计为答案，因为(v1,v2)的LCA不是u。

#include <bits/stdc++.h>
#define ll long long
using namespace std;
const int N = 2e5+10;
int n;
struct node{
	int to,w,nex;
}g[N<<1];
int head[N],cnt=0;
void add(int u,int v,int w){
	g[cnt]=node{v,w,head[u]};
	head[u]=cnt++;
}
int son[N],sz[N],dep[N];
ll dis[N],k;
void getsz(int u,int fa){
	int mx=-1;
	sz[u]=1;
	for(int i=head[u];~i;i=g[i].nex){
		int v=g[i].to;
		if(v==fa) continue;
		dep[v]=dep[u]+1;
		dis[v]=dis[u]+g[i].w;
		getsz(v,u);
		sz[u]+=sz[v]; 
		if(mx<sz[v]){
			son[u]=v;
			mx=sz[v];
		}
	}
}
unordered_map<ll,int> mp;
int temp[N],cntt=0;
int num[N],ma[N],SON;
int ans;
ll need;
void dfs1(int uu,int u,int fa){	
	temp[++cntt]=u;
	need=k+(dis[uu]<<1)-dis[u];
	if(mp[need]!=0) ans=min(ans,mp[need]+dep[u]-(dep[uu]<<1));
	for(int i=head[u];~i;i=g[i].nex){
		int v=g[i].to;
		if(v==fa||v==SON) continue;
		dfs1(uu,v,u);
	}	
}
void dfs(int u,int fa,bool keep){
	for(int i=head[u];~i;i=g[i].nex){
		int v=g[i].to;
		if(v==fa||v==son[u]) continue;
		dfs(v,u,0);
	}	
	if(son[u]!=-1){
		dfs(son[u],u,1);
		SON=son[u];
	}
	for(int i=head[u];~i;i=g[i].nex){
		int v=g[i].to;
		if(v==fa||v==son[u]) continue;
		cntt=0;
		dfs1(u,v,u);
		for(int j=1;j<=cntt;j++){
			if(mp[dis[temp[j]]]!=0)
				mp[dis[temp[j]]]=min(mp[dis[temp[j]]],dep[temp[j]]);	
			else mp[dis[temp[j]]]=dep[temp[j]];
		}
	}
	need=k+dis[u];
	if(mp[need]!=0){
		ans=min(ans,mp[need]-dep[u]);
	}
	if(mp[dis[u]]!=0)
		mp[dis[u]]=min(mp[dis[u]],dep[u]);	
	else mp[dis[u]]=dep[u];	
	SON=0;
	if(!keep) mp.clear();
}
int main(void){
	scanf("%d%lld",&n,&k);
	ans=n;
	for(int i=1;i<=n;i++)
		head[i]=son[i]=-1;
	for(int i=1;i<n;i++){
		int u,v,w;
		scanf("%d%d%d",&u,&v,&w);
		u++,v++;
		add(u,v,w);
		add(v,u,w); 
	}
	dis[1]=0;
	dep[1]=1;
	getsz(1,0);
	dfs(1,0,0);
	if(ans==n) ans=-1;
	cout<<ans;
	return 0;	
}