树链剖分分为2步,第一次求出深度,重儿子,第二次求出重链,用到了启发式的思想,即对于比较重的儿子,尽量去完整的维护它。类似于我们去合并两个堆,明显把小的堆逐个插入大的堆中会比大的往小的插更优,而这可以达到均摊O(logn)的效果。对于这个题,类似选重儿子, 我们每次尽量选择最长的路径,选出前m个就可以了。
代码:
#include <bits/stdc++.h>
#define LL long long
using namespace std;
const int maxn = 100010;
int head[maxn], Next[maxn * 2], ver[maxn * 2], tot;
LL mx_d[maxn], w[maxn], dist[maxn], d[maxn];
void add(int x, int y) {
ver[++tot] = y;
Next[tot] = head[x];
head[x] = tot;
}
void dfs1(int x, int fa = 0) {
mx_d[x] = d[x];
for (int i = head[x]; i; i = Next[i]) {
int y = ver[i];
if(y != fa) {
d[y] = d[x] + w[y];
dfs1(y, x);
mx_d[x] = max(mx_d[x], mx_d[y]);
}
}
}
void dfs2(int x, int fa, LL now_dist) {
LL mx = -1, tmp = 0;
dist[x] = now_dist;
int pos = 0;
for (int i = head[x]; i; i = Next[i]) {
int y = ver[i];
if(y == fa) continue;
if(mx_d[x] == mx_d[y]) {
pos = y;
break;
}
}
if(pos != 0) {
dfs2(pos, x, dist[x] + w[pos]);
dist[x] = 0;
}
for (int i = head[x]; i; i = Next[i]) {
int y = ver[i];
if(y == fa || y == pos) continue;
dfs2(y, x, w[y]);
}
}
int main() {
int T, n, m, x, y, kase = 0;
cin >> T;
while(T--) {
memset(head, 0, sizeof(head));
tot = 0;
scanf("%d%d", &n, &m);
for (int i = 1; i <= n; i++) {
scanf("%lld", &w[i]);
}
d[1] = dist[1] = w[1];
for (int i = 1; i < n; i++) {
scanf("%d%d", &x, &y);
add(x, y);
add(y, x);
}
dfs1(1);
dfs2(1, 0, w[1]);
LL ans = 0;
sort(dist + 1, dist + n + 1);
for (int i = n; i >= n - m + 1; i--) {
ans += dist[i];
}
printf("Case #%d: ", ++kase);
printf("%lld\n", ans);
}
}