黑白染色树

黑白染色树

题意

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

思路

树形 dp，定义 $d{p}_{i,j}$ 表示以 $i$ 为根的字数内染了 $j$ 个节点为黑色的最大收益值。

转移方程：

$$d{p}_{i,j+t}=d{p}_{i,j}+d{p}_{v,t}+(t\times (k-t)+(n-k-si{z}_v+t)\times (si{z}_v-t))\times w(i,v)$$

即所有 $v$ 子树内的黑点和所有 $v$ 子树外的黑点连边时都要经过边 $(i,v)$，白点同理。

计算出 $v$ 子树内的黑点个数和白点个数，用总黑点个数和总白点个数减去它们，相乘就是 $(i,v)$ 被经过的次数，最后乘上边权就是贡献。

注意倒序枚举 $j$ 避免后效性。

代码

#include <bits/stdc++.h>
#define int long long
using namespace std;
const int N = 2e3 + 5;
int n, k, tot, siz[N], dp[N][N];
int ver[N << 1], nxt[N << 1], head[N], edge[N << 1];
void add(int x, int y, int z) {
	ver[++ tot] = y;
	nxt[tot] = head[x];
	head[x] = tot;
	edge[tot] = z;
}
void dfs(int x, int fa) {
	siz[x] = 1;
	for (int i = head[x], y; i; i = nxt[i]) 
		if ((y = ver[i]) != fa) {
			dfs(y, x);
			for (int j = siz[x]; j >= 0; j --) 
				for (int t = siz[y]; t >= 0; t --) 
					dp[x][j + t] = max(dp[x][j + t], dp[x][j] + dp[y][t] + edge[i] * t * (k - t) + edge[i] * (n - k - siz[y] + t) * (siz[y] - t));
			siz[x] += siz[y];
		}
}
signed main() {
	cin >> n >> k;
	for (int i = 1, u, v, w; i < n; i ++) {
		cin >> u >> v >> w;
		add(u, v, w); add(v, u, w);
	}
	dfs(1, 0);
	cout << dp[1][k] << "\n";
	return 0;
}