洛谷 P3177 HAOI2015 树上染色

url: https://www.luogu.com.cn/problem/P3177

tag:
动态规划，树形DP

思路：
用 f[u][j] 表示以u为根节点的子树中，恰好有j个黑点时的收益。之后就是和分组背包类似的思路，对于每一个相连的节点，由当那个节点取k个黑点时更新过来，在不同的k值之间取一个最大值。使用 (k * (m - k) + (s[v] - k) * (n - m - s[v] + k)) * w[i] 来计算u和v之间边对答案的贡献。最后记得开long long。

代码：

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;
const int N = 2020;
int h[N], e[N * 2], ne[N * 2], w[N * 2], idx;
long long n, m;
long long s[N];
long long f[N][N];
void add(int a, int b, int c)
{
    e[idx] = b;
    w[idx] = c;
    ne[idx] = h[a];
    h[a] = idx ++;
}
void dfs(int u, int fa)
{
    s[u] = 1;f[u][0] = f[u][1] = 0;
    for (int i = h[u]; ~i; i = ne[i])
    {
        int v = e[i];
        if (v == fa) continue;
        dfs(v, u);
        s[u] += s[v];
        for (int j = min(s[u], m); j >= 0; j --)
            for (int k = 0; k <= min((long long)j, s[v]); k ++)
                if (f[u][j - k] != -1)
                {
                    long long val = (k * (m - k) + (s[v] - k) * (n - m - s[v] + k)) * w[i];
                    f[u][j] = max(f[u][j], f[u][j - k] + f[v][k] + val);
                }
    }
}
int main()
{
    memset(h, -1, sizeof h);
    memset(f, -1, sizeof f);
    cin >> n >> m;
    for (int i = 0; i < n - 1; i ++)
    {
        int a, b, c;
        cin >> a >> b >> c;
        add(a, b, c), add(b, a, c);
    }
    dfs(1, 0);
    cout << f[1][m] << endl;
    return 0;
}