携手创作，共同成长！这是我参与「掘金日新计划 · 8 月更文挑战」的第20天，点击查看活动详情

[HAOI2015] 树上染色

题目描述

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

输入格式

第一行包含两个整数 $n,k$ 。

第二到 $n$ 行每行三个正整数 $fr, to, dis$ ，表示该树中存在一条长度为 $dis$ 的边 $(fr, to)$ 。输入保证所有点之间是联通的。

输出格式

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

样例 #1

样例输入 #1

3 1
1 2 1
1 3 2

样例输出 #1

提示

对于 $100\%$ 的数据， $0 \leq n,k \leq 2000$ 。

转移方程为dp[u][i] = max( dp[u][i], dp[u][i-j] + dp[v][j] + val )

其中v为u的子节点，j为在这个子节点中选择的黑色点的个数，val为这条边的贡献

val = j*(k-j)w + (sz[v]-j)(n-k+j-sz[v])*w

其中w为这条边的边权，n为总的节点数，k为总的需要选择的黑色节点数，sz[v]为以v为根的子树的节点数量

#include <iostream>
#include <cstring>
#include <cstdio>
#include <algorithm>
#include <cstdlib>
#include <cmath>
using namespace std;
typedef long long ll;
const ll INFLL = 0x3f3f3f3f3f3f3f3f;
const int INF = 0x3f3f3f3f;
const int MAXN = 2010;
int n,k;
struct Tree {
    int head[MAXN], nxt[MAXN<<1], to[MAXN<<1], w[MAXN<<1], idx;
    Tree() { idx = 0; memset( head, -1, sizeof(head) ); }
    void addedge( int u, int v, int val ) {
        to[idx] = v; nxt[idx] = head[u]; w[idx] = val; head[u] = idx; ++idx;
        to[idx] = u; nxt[idx] = head[v]; w[idx] = val; head[v] = idx; ++idx;
    }
}tree;
ll dp[MAXN][MAXN]; int sz[MAXN];
void dfs( int u, int fa ) {
    sz[u] = 1; memset( dp[u], -1, sizeof(dp[u]) ); dp[u][0] = dp[u][1] = 0;
    for( int e = tree.head[u]; ~e; e = tree.nxt[e] ) {
        int v = tree.to[e]; if( v == fa ) continue; dfs(v,u); sz[u] += sz[v];
    }
    for( int e = tree.head[u]; ~e; e = tree.nxt[e] ) {
        int v = tree.to[e]; if( v == fa ) continue; int w = tree.w[e];
        for( int i = min(k,sz[u]); i >= 0; --i ) for( int j = 0; j <= min(i,sz[v]); ++j ) if( ~dp[u][i-j] ) {
            ll val = (ll)j*(k-j)*w + (ll)(sz[v]-j)*(n-k+j-sz[v])*w;
            dp[u][i] = max( dp[u][i], dp[u][i-j] + dp[v][j] + val );
        }
    }
}
int main() {
    scanf( "%d%d", &n, &k );
    for( int i = 0; i < n-1; ++i ) {
        int u,v,w; scanf( "%d%d%d", &u, &v, &w ); tree.addedge(u,v,w);
    } dfs(1,0); printf( "%lld\n", dp[1][k] );
    return 0;
}

【动态规划】树上染色