P3177 [HAOI2015] 树上染色

[P3177 HAOI2015] 树上染色 - 洛谷 | 计算机科学教育新生态 (luogu.com.cn)

P3177 [HAOI2015] 树上染色

题目大意

有一棵 \(n\) 个点的树，你可以在上面把 \(k\) 个点染成黑色，收益为黑点两两之间的距离和加上白点两两之间的距离和

求最大收益。

思路

树上背包

这种题肯定找边的贡献

设 \(f_{u , i}\) 为以 \(u\) 为根的子树内染了 \(i\) 个黑点的最大收益

\(sz_u\) 为以 \(u\) 为根的子树的节点个数

一条连接 \(u , v\) 的边， \(v\) 子树上有 \(j\) 个黑点，那么这条边的贡献为：

\[val = j * (k - j) * e[i].w + (sz_v - j) * (n - k - sz_v + j) * e[i].w \]

显然：

\[f_{u , i} = \max_{v \in son(u)} f_{v , j} +f_{u , i - j} + val (0\le j\le i \le k ) \]

实现时要倒序枚举 \(i\) 防止重复选， \(j = 0\) 的情况要提前转移一下

code

#include <bits/stdc++.h>
#define fu(x , y , z) for(int x = y ; x <= z ; x ++)
#define fd(x , y , z) for(int x = y ; x >= z ; x --)
#define LL long long
using namespace std;
const int N = 2005;
int n , hd[N] , cnt , fa[N] , sz[N] , K;
LL f[N][N];
struct E {
    int to , nt;
    LL w;
} e[N << 1];
void add (int x , int y , LL z) { e[++cnt].to = y , e[cnt].nt = hd[x] , hd[x] = cnt , e[cnt].w = z; }
void dfs (int x) {
    int y;
    sz[x] = 1;
    f[x][0] = f[x][1] = 0;
    for (int i = hd[x] ; i ; i = e[i].nt) {
        y = e[i].to;
        if (y == fa[x]) continue;
        fa[y] = x;
        dfs (y);
        sz[x] += sz[y];
        fd (j , min (K , sz[x]) , 0) {
            if (f[x][j] != -1)
                f[x][j] += f[y][0] + (n - K - sz[y]) * sz[y] * e[i].w;
            fu (k , max(1 , j + sz[y] - sz[x]) , min (j , sz[y])) {
                if (f[x][j - k] == -1) continue;
                f[x][j] = max (f[x][j] , f[x][j - k] + f[y][k] + e[i].w * (K - k) * k + e[i].w * (n - K - sz[y] + k) * (sz[y] - k)); 
            } 
        }
    }
}
int main () {
    int u , v;
    LL w;
    scanf ("%d%d" , &n , &K);
    fu (i , 1 , n - 1) {
        scanf ("%d%d%lld" , &u , &v , &w);
        add (u , v , w) , add (v , u , w); 
    }
    memset (f , -1 , sizeof (f));
    dfs (1);
    printf ("%lld" , f[1][K]);
    return 0;
}