黑白染色树
题意
有一棵点数为 \(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;
}