【每日一题】[HAOI2015]树上染色（树形dp）

Description

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

Solution

$n, k \leq 2000$ , 考虑二维树形 $dp$
令 $dp[i][j]$ 表示以 $i$ 为根，选取 $j$ 个作为黑色节点的收益
考虑每天边权的贡献，类似于树上两点的距离，对于一条边，它连接两端的相同颜色节点都会产生贡献
假设边的左边有 $black_1$ 个黑节点， $while_1$ 个白节点，右边有 $black_2$ 个黑节点， $while_2$ 个白节点，边权为 $w$
那么贡献为 $w \cdot black_1 \cdot black_2 + w \cdot white_1 \cdot white_2$
但是显然没有用到 $k$ 的条件
考虑在 $dfs$ 的时候记录子树大小 $sz[i]$
枚举到根节点 $u$ 和子节点 $v$ 的黑色节点个数分别为 $i, j$
那么有 $dp[u][i + j] = dp[v][j] + dp[u][i] + val$
其中 $val$ 为连接两点的边权的贡献
即 $w \cdot j \cdot (k - j) + w \cdot (sz[v] - j) \cdot ((n - k) - (sz[v] - j))$
考虑消除后效性，类似于背包问题，我们倒叙枚举即可。

#include<bits/stdc++.h>

typedef long long ll;
using namespace std;
#define emplace_back push_back

const int N = 1e6 + 5;
const int mod = 998244353;

vector<pair<int, int> > G[N];
ll dp[2005][2005], n, k;
ll sz[2005];

ll cal(int u, int v, int w, int i, int j) {
  return dp[u][i] + dp[v][j] + 1LL * w * j * (k - j) + w * (n - k - (sz[v] - j)) * (sz[v] - j);
}
void dfs(int u, int par) {
  sz[u] = 1;
  for (auto &it : G[u]) {
    int v = it.first;
    if (v == par) continue;
    dfs(v, u);
    for (int i = sz[u]; i >= 0; i--) {
      for (int j = sz[v]; j >= 0; j--) {
        dp[u][i + j] = max(dp[u][i + j], cal(u, v, it.second, i, j));
      }
    }
    sz[u] += sz[v];
  }
}
void solve() {
  cin >> n >> k;
  for (int i = 0; i < n - 1; i++) {
    int u, v, w; cin >> u >> v >> w;
    G[u].emplace_back({v, w});
    G[v].emplace_back({u, w});
  }
  dfs(1, 0);
  cout << dp[1][k] << '\n';
}
int main() {
  ios::sync_with_stdio(false), cin.tie(0),  cout.tie(0);
  int T = 1; //cin >> T;
  while(T--) {
    solve();
  }
  return 0;
}