[HAOI 2015]树上染色

Description

题库链接

给出一棵 $n$ 个节点的树，边有权值。让你将树上 $k$ 个点染黑，剩余 $n-k$ 个点染白。染色后记一种染色方案的价值为黑点间两两距离和以及白点间两两距离和。求最大价值。

$0\leq k\leq n\leq 2000$

Solution

定义状态 $f_{u,i}$ 表示以 $u$ 为根的子树中选出了 $i$ 个黑点的子树中最大价值。

转移的话就是考虑当前节点和枚举的儿子间的边被计算了几次。

树上背包复杂为 $O(n^2)$ 。

Code

#include <bits/stdc++.h>
#define ll long long
using namespace std;
const int N = 2000+5;
 
int n, k, u, v, size[N]; ll f[N][N], c;
struct tt {int to, next; ll cost; }edge[N<<1];
int path[N], top;
 
void add(int u, int v, ll cost) {edge[++top] = (tt){v, path[u], cost}; path[u] = top; }
void dfs(int u, int fa) {
    size[u] = 1; ll c = 0;
    for (int i = path[u], v; i; i = edge[i].next) {
        if ((v = edge[i].to) == fa) continue;
        dfs(v, u); c = edge[i].cost;
        for (int j = min(size[u], k); j >= 0; j--)
            for (int p = min(k-j, size[v]); p >= 0; p--)
                f[u][j+p] = max(f[u][j+p], f[u][j]+f[v][p]+1ll*p*(k-p)*c+1ll*(n-k-size[v]+p)*(size[v]-p)*c);
        size[u] += size[v];
    }
}
void work() {
    scanf("%d%d", &n, &k);
    for (int i = 1; i < n; i++) {
        scanf("%d%d%lld", &u, &v, &c); add(u, v, c), add(v, u, c);
    }
    dfs(1, 0);
    printf("%lld\n", f[1][k]);
}
int main() {work(); return 0; }