P3177 [HAOI2015]树上染色

Description

有一棵点数为 N 的树，树边有边权。给你一个在 0~ N 之内的正整数 K ，你要在这棵树中选择 K个点，将其染成黑色，并将其他 的N-K个点染成白色 。 将所有点染色后，你会获得黑点两两之间的距离加上白点两两之间的距离的和的受益。问受益最大值是多少。

Solution

\(f_{i,j}\)表示\(i\)个点选择\(j\)个黑点的最大收益.

Code

#include <cstdio>
#include <cstring>
#include <iostream>
#define N 2005
using namespace std;
#define int long long

struct Edge {
    int v, dis, nxt;
} e[N << 1];
int head[N], tot;

void AddEdge(int u, int v, int d) {
    e[++tot] = (Edge) {v, d, head[u]}; head[u] = tot;
    e[++tot] = (Edge) {u, d, head[v]}; head[v] = tot;
}
int son[N], f[N][N];
int n, k;

int find(int u, int fa) {
    son[u] = 1; f[u][0] = f[u][1] = 0;
    for (int i = head[u]; i; i = e[i].nxt) {
        if (e[i].v == fa)continue;
        son[u] += find(e[i].v, u);
        for (int ss = min(son[u], k); ss >= 0; --ss)
            for (int j = 0; j <= min(son[e[i].v], ss); ++j) {
                f[u][ss] = max(f[u][ss], f[u][ss - j] + f[e[i].v][j] +
                    (j * (k - j) + (son[e[i].v] - j) * (n - k + j - son[e[i].v])) * e[i].dis);
            }
    }
    return son[u];
}

main() {
    int u, v, d;
    scanf("%lld%lld", &n, &k);
    for (int i = 1; i < n; ++i) {
        scanf("%lld%lld%lld", &u, &v, &d);
        AddEdge(u, v, d);
    }
    for (int i = 1; i <= n; ++i)
        for (int j = 1; j <= k; ++j)
            f[i][j] = -123456789;
    find(1, 0);
    printf("%lld\n", f[1][k]);
    return 0;
}