火山哥周游世界 (树形dp)

题意：一共有n个国家，标号1~n，每条边都有一定的时间花费，任意两个国家之间两两可达。
火山哥一共决定去K个国家。现在他想要知道:如果他从第 i 个国家出发，经过这 K 个国家的最短时间是多少？

输入：
3 3 (n，k)
1 2 1
2 3 4
1
2
3
输出：
5
6
5

输入：
5 2
1 3 1
2 3 2
4 3 3
5 1 4
5
1
输出：
4
7
5
8
4

解题思路：
出发点与根节点与所有关键点构成的子树边权总和乘2-出发点到最远关键点的距离

#include<bits/stdc++.h>
using namespace std;
#define maxn 1000010
typedef long long ll;
struct Edge{
	int to,w,nxt;
}edge[maxn];
int n,k,tot=0,vis[maxn],siz[maxn],head[maxn];
ll dis1[maxn],dis2[maxn],sum[maxn];
void add(int u,int v,int w){
	edge[++tot].to=v; edge[tot].w=w; edge[tot].nxt=head[u]; head[u]=tot;
	edge[++tot].to=u; edge[tot].w=w; edge[tot].nxt=head[v]; head[v]=tot;
}
void cmp(ll x,int u){
    if(x>=dis1[u]) dis2[u]=dis1[u],dis1[u]=x;
    else if(x>dis2[u]) dis2[u]=x;
}
void dfs1(int u,int fa){
	if(vis[u]) siz[u]=1;
	sum[u]=dis1[u]=dis2[u]=0;
    for(int i=head[u];~i;i=edge[i].nxt){
    	int to=edge[i].to,w=edge[i].w;
    	if(fa==to) continue;
    	dfs1(to,u);
    	siz[u]+=siz[to];
        sum[u]+=sum[to]+(siz[to]!=0)*2*w;
        cmp(dis1[to]+(siz[to]!=0)*w,u);
	}
}
void dfs2(int u,int fa){
	for(int i=head[u];~i;i=edge[i].nxt){
    	int to=edge[i].to,w=edge[i].w;
    	if(fa==to) continue;
    	sum[to]=sum[u]-2*w*(siz[to]!=0)+(n-siz[to]!=0)*2*w;
    	if(siz[to]!=n){
            if(dis1[to]+w==dis1[u]) cmp(dis2[u]+w,to);
            else cmp(dis1[u]+w,to);
        }
    	dfs2(to,u);
	}
}
int main(){
	memset(vis,0,sizeof(vis));
	memset(head,-1,sizeof(head));
	scanf("%d%d",&n,&k);
	for(int i=1;i<n;i++){
		int u,v,w; scanf("%d%d%d",&u,&v,&w);
		add(u,v,w);
	}
	for(int i=1;i<=k;i++){
		int x; scanf("%d",&x); 
		vis[x]=1;
	}
	dfs1(1,0); dfs2(1,0);
	for(int i=1;i<=n;i++) printf("%lld\n",sum[i]-dis1[i]);
	return 0;
}