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

Description

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

Solution

, 考虑二维树形
令 表示以 为根，选取个作为黑色节点的收益
考虑每天边权的贡献，类似于树上两点的距离，对于一条边，它连接两端的相同颜色节点都会产生贡献
假设边的左边有 个黑节点， 个白节点，右边有 个黑节点， 个白节点，边权为
那么贡献为
但是显然没有用到 的条件
考虑在 的时候记录子树大小
枚举到根节点 和子节点 的黑色节点个数分别为
那么有
其中 为连接两点的边权的贡献
即
考虑消除后效性，类似于背包问题，我们倒叙枚举即可。

#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;
}