[HAOI2015]树上染色

[HAOI2015]树上染色

题目链接

题意：

给你一颗大小为$n$的树，然后将树上的$k$个点染成黑色， $n-k$个点染成白色。问黑色两两之间的距离，与白色两两之间的距离和最大是多少？

题解：

这题写了一晚上，发现自己dp方程就设错了。看了题解才知道怎么写。网上题解很多，但是很多都忽略一个很重要的细节。下面我详细说一下。

设 $dp[u][j]$表示$u$为节点的子树中选$j$个黑点， 黑点两两之间的距离与 剩下的$k-j$个子树外的黑的点经过 子树$u$黑点的边（属于子树u里面的边）产生的贡献， 同理，再加上 子树中剩下白色点两两距离， 与子树外白色点到子树内的白色点边的和的贡献最大。

有点难理解， 网上很多题解 $dp$表示的含义就说错了。建议多看几遍。

怎么转移？

设$sz[u]$表示这个与$u$根节点的子树大小。

$to$ 为$u$的儿子节点， $w$为$u$到$to$的距离。

方程转移时考虑边产生的贡献， $w$这条件产生的贡献为：

节点$to$ 的黑点个数 $\times$ $(k-节点to的黑的个数)\times w$ + 节点$to$白点个数$\times$节点$to$以外的白点个数 $\times$ w

那么：

设$dp[u][i+j]$为再$u$为节点的子树中选$i$个黑点， 再$u$的儿子节点为子树中选$j$个黑点。

$dp[u][i+j]=max(dp[u][i+j],dp[u][i]+dp[to][j]+val)$

其中 $val=(j\ast (k-j)+(sz[to]-j)\ast (n-k-sz[to]+j))\ast w$

如果想不通为啥$val$的值是这样计算的，建议再读一边我上面的$dp$方程的定义。

直接枚举$i,j$就可了， 根据子树的大小枚举可以将复杂度控制再$n^2$， 因为总共将每队节点都枚举，总共有$n^2$对

时间复杂度为$o(n+n^2)$， 主意用个滚动数组记录。

代码：

#include<bits/stdc++.h>
using namespace std;
const int N = 2007;

typedef long long ll;
vector<pair<int, int> > g[N];
ll n, k;
ll dp[N][N], temp[N], sz[N];

void dfs(int u, int fa, ll cost) {
    sz[u] = 1;
    for (auto it: g[u]) {
        int to = it.first;
        ll w = it.second;
        if (to == fa) continue;
        dfs(to, u, w);

        for (ll i = 0; i <= sz[u]; i++) {
            for (ll j = 0; j <= sz[to]; j++) {
                if (i + j > k) break;
                ll val = j * (k - j) * w + (sz[to] - j) * (n - k - sz[to] + j) * w;
                temp[i + j] = max(temp[i + j], dp[u][i] + dp[to][j] + val);

            }
        }

        for (int i = 0; i <= sz[u]; i++) {
            for (int j = 0; j <= sz[to]; j++) {
                if (i + j > k) break;
                dp[u][i + j] = max(dp[u][i + j], temp[i + j]);
                temp[i + j]= 0;
            }
        }
        sz[u] += sz[to];
    }
    
}

int main() {
    scanf("%lld %lld", &n, &k);
    for (int i = 1; i < n; i++) {
        int u, v, w; scanf("%d %d %d", &u, &v, &w);
        g[u].push_back({v, w});
        g[v].push_back({u, w});
    }
    dfs(1, 0, 0);
    cout << dp[1][k] << endl;
}