树链剖分分为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 ", ans); } }