9. 景区导游-树上前缀和&&最近公共祖先

9.景区导游 - 蓝桥云课 (lanqiao.cn)

6 4
1 2 1
1 3 1
3 4 2
3 5 2
4 6 3
2 6 5 1

#include <bits/stdc++.h>
#define int long long
#define endl '\n'
using namespace std;
//tu动态数组用来存图,相当于一条绳子上面有好几个结，每个结都各自连了一条链子
//vector <pair<to,value>>tu[maxn] maxn->节点最大数量，to->联通的节点,value->权值
vector <pair<int,int>>tu[100002];
//dep存节点的深度，把图看成树
//fa[u][i]存u结点往上跳2^i次幂后是哪个节点
//s存边前缀和，以游览路线的第一个景点为根节点，把其他每一个节点到根节点的权值和存到s数组里
int dep[100002],fa[100002][19],s[100002];
//树上前缀和(在最近公共祖先的板子上加以改动就行)+最近公共祖先
//dfs直接套模板
void dfs(int u,int father){//也可以说是预处理节点信息
	dep[u]=dep[father]+1;
	fa[u][0]=father;
	for(int i=1;i<=18;i++){
		fa[u][i]=fa[fa[u][i-1]][i-1];
		if(!fa[u][i]) break;
	}
	for(int i=0;i<tu[u].size();i++){
//		这里是和父节点判断
		if(tu[u][i].first==father) continue;
//		下面这一行不是模板里的，不是求最近公共祖先的模板
//		而是求树上前缀和的(递推就求出来了)
//		该点到根节点的前缀和=其父节点的前缀和+父节点到该点的权值
		s[tu[u][i].first]=s[u]+tu[u][i].second;
		dfs(tu[u][i].first,u);
	}
}
//lca函数能返回其两点的最近公共祖先
//lca一点变化都没有直接套模板
int lca(int u,int v){
	if(dep[u]<dep[v]) swap(u,v);
	for(int i=18;i>=0;i--){
//		别忘了这里是判断深度dep
		if(dep[fa[u][i]]>=dep[v])	u=fa[u][i];
	}
	if(u==v) return u;
	for(int i=18;i>=0;i--){
		if(fa[u][i]!=fa[v][i]){
			u=fa[u][i];
			v=fa[v][i];
		}
	}
	return fa[v][0];
}
//计算u节点到v节点的权值
//因为是边前缀和，所以两个节点之间的前缀和=两点的前缀和相加-两倍的其最近公共祖先的前缀和
int cal(int u,int v){
	return s[u]+s[v]-2*s[lca(u,v)];
}
void solve(){
	int n,k,shunxu[100002],sum=0;
	cin>>n>>k;
//	存图
	for(int i=0;i<n-1;i++){
		int u,v,t;
		cin>>u>>v>>t;
		tu[u].push_back({v,t});
		tu[v].push_back({u,t});
	}
//	存游览线路的顺序
	for(int i=0;i<k;i++){
		cin>>shunxu[i];
	}
//	遍历一遍，好把前缀和啥的都存起来了
//	记住，这里就调用了一次，然后他在递归
	dfs(shunxu[0],0);
//	把总线路的权值都加一块存起来
	for(int i=0;i<k-1;i++){
		sum+=cal(shunxu[i],shunxu[i+1]);
	}
//	依次计算跳第i个景点后的权值和
	for(int i=0;i<k;i++){
//		如果是跳过第一个景点，那只用把第一条线减掉
		if(i==0){
			cout<<sum-cal(shunxu[0],shunxu[1])<<" ";
		}
//		如果是跳过最后一个景点，那只用把最后一条线减掉
		else if(i==k-1){
			cout<<sum-cal(shunxu[k-2],shunxu[k-1])<<" ";
		}
//		如果是跳过中间景点，那需要把该节点左右两边的线减掉，再加上该点左右两个景点之间的路径
		else{
			cout<<sum-cal(shunxu[i-1],shunxu[i])-cal(shunxu[i],shunxu[i+1])+cal(shunxu[i-1],shunxu[i+1])<<" ";
		}
	}
}
signed main(){
	ios::sync_with_stdio(0);
	cin.tie(0);
	cout.tie(0);
	solve();
	return 0;
}