BZOJ3252 攻略

Problem

给定一棵包含 \(n\) 个节点的有根树。定义一条路径为从根节点到一个叶子节点的简单路径，你需要选出 \(k\) 条这样的路径。

每个节点有一个权值 \(w\)，你的收益就是这 \(k\) 条路径所覆盖到的点的权值和。（覆盖多次不能累加权值）

现在想要最大化收益，请给出最大的收益值。

\(1 \le n \le 2 \times 10^5，1 \le w_i \le 2^{31}-1\)

Input

第一行两个正整数 \(n\)，\(k\)。

第二行 n 个正整数，表示每个节点的权值。

接下来 \(n-1\) 行，每行两个整数 \(x\)，\(y\)，表示从 \(x\) 向 \(y\) 连一条有向边，即 \(x\) 为 \(y\) 的父亲。

数据保证节点 \(1\) 为根节点

Output

输出一个整数表示最大的收益。

Sample

Input 1

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

Output 1

10

Solution

考虑将原树类似树链剖分那样切开，但这里不是根据子树大小来断开，而是比较权值大小。（注意要先将链上节点的权值求和再比较）

然后对于每一条链，权值求和塞进堆中，取出最大的前 \(k\) 个统计答案即可。

证明一下正确性：

如果一个节点与父亲断开被记为轻链，那这个父亲节点的重儿子上链的权值和一定是大于该点链上权值和的，选择时优先重儿子一定是正确的。

因此，假如一个轻儿子被选择，那其父亲节点的重儿子那条链一定被选择过了，该点到根节点的路径上的节点权值也是被统计过了的。

代码：

#include <algorithm>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <queue>

using namespace std;

const int kmax = 2e5 + 3;

struct E {
  int p, y;
} e[kmax << 1];

int n, k;
long long w[kmax], res;
int h[kmax], ec;
int son[kmax];
priority_queue<long long> q;

void Addedge(int x, int y) {
  e[++ec] = {h[x], y};
  h[x] = ec;
}

void Dfs(int x, int fa) {
  for (int i = h[x]; i; i = e[i].p) {
    int y = e[i].y;
    if (y == fa) continue;
    Dfs(y, x);
    if (w[y] > w[son[x]]) {
      son[x] = y;
    }
  }
  w[x] += w[son[x]];
}

void Dfss(int x, int fa, int tp) {
  if (x == tp) q.push(w[x]);
  if (son[x]) Dfss(son[x], x, tp);
  for (int i = h[x]; i; i = e[i].p) {
    int y = e[i].y;
    if (y == fa || y == son[x]) continue;
    Dfss(y, x, y);
  }
}

int main() {
  freopen("game.in", "r", stdin);
  freopen("game.out", "w", stdout);
  scanf("%d%d", &n, &k);
  for (int i = 1; i <= n; i++) {
    scanf("%lld", &w[i]);
  }
  for (int i = 1, x, y; i < n; i++) {
    scanf("%d%d", &x, &y);
    Addedge(x, y);
  }
  Dfs(1, 0);
  Dfss(1, 0, 1);
  for (int i = 1; i <= k && !q.empty(); i++, q.pop()) {
    res += q.top();
    //cout << q.top() << '\n';
  }
  printf("%lld\n", res);
  return 0;
}