P6419 [COCI2014-2015#1] Kamp

P6419 COCI2014-2015#1 Kamp
营地聚会

点击查看代码

#include <iostream>
#include <stdio.h>
#include <string.h>
#include <algorithm>
#include <utility>
#include <array>
#include <queue>

using namespace std;

typedef long long LL;

const int N = 5e5 + 5, M = 1e6 + 5;

int n, k;
int h[N], e[M], w[M], nxt[M], idx;
bool st[N]; // 是否有人
LL g[N]; // g[u] 表示以 u 为根节点将住在以 u 为根的子树中的人都送到家再回到 u 的最小步数
int sz[N];// sz[u] 为住在以 u 为根的子树的人数, g[u] = ∑{g[v]+2*w[u,v] , sz[v]>0}
LL d1[N], d2[N]; int p1[N]; // 最长链, 次长链, ...
LL f[N]; // f[u] 表示以 u 为根节点将所有人送回家再回到 u 的最小步数

void add(int a, int b, int c) {
	e[++ idx] = b, w[idx] = c, nxt[idx] = h[a], h[a] = idx;
}

void dfs1(int u, int fa) {
	if(st[u]) sz[u] = 1;
	for(int i = h[u]; i; i = nxt[i]) {
		int v = e[i];
		if(v == fa) continue;
		dfs1(v, u);
		sz[u] += sz[v];
		if(sz[v]) {
			g[u] += g[v] + 2 * w[i];
			int d = d1[v] + w[i];
			if(d >= d1[u]) d2[u] = d1[u], d1[u] = d, p1[u] = v;
			else if(d > d2[u]) d2[u] = d;
		}
	}
}

void dfs2(int u, int fa) {
	for(int i = h[u]; i; i = nxt[i]) {
		int v = e[i];
		if(v == fa) continue;
		if(!sz[v]) { // 子树没人
			f[v] = f[u] + 2 * w[i];
			d1[v] = d1[u] + w[i];
		} else if(sz[v] == k) {
			f[v] = g[v]; // 只有子树有人
		} else {
			f[v] = f[u];
			/* */if(d1[u] + w[i] > d1[v] && p1[u] != v)
				d2[v] = d1[v], d1[v] = d1[u] + w[i], p1[v] = u;
//			else if(d1[u] + w[i] > d1[v] && p1[u] == v); // 不能更新
			else if(d2[u] + w[i] > d1[v])
				d2[v] = d1[v], d1[v] = d2[u] + w[i], p1[v] = 1;
			else if(d1[u] + w[i] > d2[v] && p1[u] != v)
				d2[v] = d1[u] + w[i];
//			else if(d1[u] + w[i] > d2[v] && p1[u] == v); // 不能更新
			else if(d2[u] + w[i] > d2[v])
				d2[v] = d2[u] + w[i];
		}
		dfs2(v, u);
	}
}

int main() {
	scanf("%d%d", &n, &k);
	for(int i = 1, a, b, c; i < n; i ++) scanf("%d%d%d", &a, &b, &c), add(a, b, c), add(b, a, c);
	for(int i = 1, x; i <= k; i ++) scanf("%d", &x), st[x] = true;
	dfs1(1, -1), f[1] = g[1], dfs2(1, -1);
	for(int i = 1; i <= n; i ++) printf("%lld\n", f[i] - d1[i]);
	return 0;
}