黑白染色树

黑白染色树

题意

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

思路

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

转移方程：

\[dp_{i,j+t}=dp_{i,j}+dp_{v,t}+(t\times(k-t)+(n-k-siz_v+t)\times(siz_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;
}