[HAOI2015] 树上染色

树形 \(DP\) , 一道有点恶心的树上背包.

设状态 \(f[x][i]\) 表示以 \(x\) 为根节点的子树中选 \(i\) 个节点染黑的最大价值.

考虑转移, 枚举以 \(x\) 为根节点的子树中选 \(j\) 个黑点, \(x\) 的子节点 \(y\) 选 \(k\) 个黑点, 那么 \(dis(i, j)\) 就被计算了 \(k * (m - k) + (son[y] - k) * (n - son[y] - m + k)\) 次, 也就是 \(子树内的黑点数 * 子树外的黑点数 + 子树内的白点数 * 子树外的白点数\) , 那么 \(f[x][j] = \max(f[x][j], f[y][k] + f[x][j - k] + (k * (m - k) + (son[y] - k) * (n - son[y] - m + k)) * z\) .

为什么说这个题恶心呢, 因为它有一个极端数据...也就是 \(Subtask\ \#1\) ...这个点是链, 跑得巨慢, 使得我们的代码效率远达不到树上背包在随机数据上的效率(所以要开 \(O2\) bushi) .

\(code:\)

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
int read() {
    int x = 0, f = 1;
    char ch = getchar();
    while (!isdigit(ch)) {
        if (ch == '-') f = -1;
        ch = getchar();
    }
    while (isdigit(ch)) {
        x = (x << 1) + (x << 3) + (ch ^ 48);
        ch = getchar();
    }
    return x * f;
}
const int N = 2005, M = 4005;
int n, m, son[N];
int tot, to[M], dis[M], nxt[M], head[N];
ll f[N][N];
void add(int u, int v, int w) {
    to[++tot] = v, dis[tot] = w, nxt[tot] = head[u], head[u] = tot;
}
void dfs(int x, int fa) {
    son[x] = 1; f[x][0] = f[x][1] = 0;
    for (int i = head[x]; i; i = nxt[i]) {
        int y = to[i], z = dis[i];
        if (y == fa) continue;
        dfs(y, x);
        son[x] += son[y];
        for (int j = min(son[x], m); j >= 0; j--) {
            if (f[x][j] != -1) f[x][j] += f[y][0] + (ll)son[y] * (n - son[y] - m) * z;
            for (int k = min(son[y], j); k > 0; k--) {
                if (f[x][j - k] == -1) continue;
                ll num = (ll)k * (m - k) + (ll)(son[y] - k) * (n - son[y] - m + k);
                f[x][j] = max(f[x][j], f[y][k] + f[x][j - k] + num * z);
            }
        }
    }
}
int main() {
    memset(f, -1, sizeof(f));
    n = read(), m = read();
    for (int i = 1; i < n; i++) {
        int u = read(), v = read(), w = read();
        add(u, v, w); add(v, u, w);
    }
    dfs(1, 0);
    printf("%lld", f[1][m]);
    return 0;
}