[HAOI2015] 树上染色

GitHub跳转原题关系图返回列表

设 dp[u][j] 为 u 子树选 j 个黑点的最大收益，把同色点对距离和拆成每条边两侧黑点对与白点对数量乘边权的贡献来转移。

OJ: luogu

题目 ID: P3177

难度:提高+/省选-

标签:树形DP动态规划树推导

日期: 2026-06-21 04:38

目录

  1. 题意
  2. 思路
    1. DP 转移方程
  3. 代码
  4. 复杂度
  5. 总结
  6. 一图流解析

题意

给一棵带边权的树，恰好选 k 个点染成黑色，其余点染成白色。

收益定义为：

  • 所有黑点两两之间的距离和
  • 加上所有白点两两之间的距离和

要求最大化这个收益。

思路

先看一个只适合小数据验证的暴力：

cpp
#include <bits/stdc++.h>
using namespace std;

const int MAXN = 25;

struct Edge {
    int to;
    int w;
};

int n, k_need;
vector<Edge> g[MAXN];
int dista[MAXN][MAXN];
int choose_black[MAXN]; // choose_black[i] = 0/1，表示第 i 个点染白/染黑
long long ans;

void dfs_dist(int start, int u, int fa, int d) {
    dista[start][u] = d;
    for (size_t i = 0; i < g[u].size(); i++) {
        int v = g[u][i].to;
        int w = g[u][i].w;
        if (v == fa) {
            continue;
        }
        dfs_dist(start, v, u, d + w);
    }
}

int calc_black_count() {
    int cnt = 0;
    for (int i = 1; i <= n; i++) {
        if (choose_black[i] == 1) cnt++;
    }
    return cnt;
}

bool check() {
    return calc_black_count() == k_need;
}

long long calc_answer() {
    long long sum = 0;
    for (int i = 1; i <= n; i++) {
        for (int j = i + 1; j <= n; j++) {
            if (choose_black[i] == choose_black[j]) {
                sum += dista[i][j];
            }
        }
    }
    return sum;
}

void dfs_choose(int dep) {
    if (dep == n + 1) {
        if (check()) {
            long long value = calc_answer();
            if (ans < value) ans = value;
        }
        return;
    }

    // 第 dep 个点的 01 选择：0 染白，1 染黑。
    for (int i = 0; i <= 1; i++) {
        choose_black[dep] = i;
        dfs_choose(dep + 1);
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    // brute.cpp：枚举哪些点染成黑色，直接统计黑黑点对和白白点对的距离和。
    cin >> n >> k_need;
    for (int i = 1; i <= n; i++) {
        g[i].clear();
    }
    for (int i = 1; i < n; i++) {
        int u, v, w;
        cin >> u >> v >> w;
        g[u].push_back({v, w});
        g[v].push_back({u, w});
    }

    for (int i = 1; i <= n; i++) {
        dfs_dist(i, i, 0, 0);
    }

    ans = 0;
    dfs_choose(1);
    cout << ans << '\n';
    return 0;
}

brute.cpp 把每个点看成一个 01 选择：choose_black[i] = 0/1 表示染白或染黑。递归先生成完整选择，叶子节点再检查是否恰好染黑 k 个点，并按点对统计答案。 这个做法正确，但复杂度是指数级。

正解的关键是把“点对距离和”拆成“边贡献”。

考虑一条父子边 u-v，设 v 是儿子，且在 v 子树中选了 b 个黑点。

那么：

  • v 子树外有 k-b 个黑点
  • 跨过这条边的黑点对数量是 b * (k-b)

所以这条边对黑点点对的贡献是：

b * (k-b) * w

同理，v 子树中白点数是 sz[v] - b，子树外白点数是 (n-k) - (sz[v]-b)， 于是白点贡献就是：

(sz[v]-b) * ((n-k) - (sz[v]-b)) * w

因此可以做树形 DP：

  • dp[u][j] 表示 u 子树里选 j 个黑点时的最大收益

合并儿子时做背包即可。

DP 转移方程

合并儿子 v 时，假设原来在已处理部分选了 a 个黑点，在 v 子树选了 b 个黑点。 这条边的新增贡献为：

add=b(k−b)w+(sz[v]−b)((n−k)−(sz[v]−b))w add=b(k-b)w + (sz[v]-b)((n-k)-(sz[v]-b))w

因此背包合并为：

dp[u][a+b]=max⁡(dp[u][a+b], dp[u][a]+dp[v][b]+add) dp[u][a+b]=\max(dp[u][a+b],\ dp[u][a]+dp[v][b]+add)

下面这张图展示了一条边的贡献来源：

graph G {
  rankdir=LR;
  L [label="v 子树"];
  R [label="子树外"];
  L -- R [label="w"];
}

只要一对同色点分居这条边两侧，它们的距离里就会包含这条边一次。 所以我们只需要统计“跨边的同色点对数”，不需要显式计算所有点对距离。

代码

cpp
#include <bits/stdc++.h>
using namespace std;

const int MAXN = 2005;

struct Edge {
    int to;
    int w;
};

int n, k_need;
vector<Edge> g[MAXN];
int sz[MAXN];
long long dp[MAXN][MAXN];

void dfs(int u, int fa) {
    sz[u] = 1;
    dp[u][0] = 0;
    dp[u][1] = 0;

    for (size_t i = 0; i < g[u].size(); i++) {
        int v = g[u][i].to;
        int w = g[u][i].w;
        if (v == fa) {
            continue;
        }

        dfs(v, u);

        static long long tmp[MAXN];
        int upper = min(k_need, sz[u] + sz[v]);
        for (int j = 0; j <= upper; j++) {
            tmp[j] = -(1LL << 60);
        }

        for (int a = 0; a <= min(k_need, sz[u]); a++) {
            for (int b = 0; b <= min(k_need - a, sz[v]); b++) {
                long long add = 0;

                // 黑点对：一边在 v 子树，一边不在 v 子树。
                add += 1LL * b * (k_need - b) * w;

                // 白点对：一边在 v 子树，一边不在 v 子树。
                int white_in_v = sz[v] - b;
                int white_out_v = (n - k_need) - white_in_v;
                add += 1LL * white_in_v * white_out_v * w;

                long long cand = dp[u][a] + dp[v][b] + add;
                if (cand > tmp[a + b]) {
                    tmp[a + b] = cand;
                }
            }
        }

        sz[u] += sz[v];
        for (int j = 0; j <= min(k_need, sz[u]); j++) {
            dp[u][j] = tmp[j];
        }
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    cin >> n >> k_need;
    for (int i = 1; i <= n; i++) {
        g[i].clear();
        for (int j = 0; j <= k_need; j++) {
            dp[i][j] = -(1LL << 60);
        }
    }

    for (int i = 1; i < n; i++) {
        int u, v, w;
        cin >> u >> v >> w;
        g[u].push_back({v, w});
        g[v].push_back({u, w});
    }

    dfs(1, 0);
    cout << dp[1][k_need] << '\n';
    return 0;
}

复杂度

树形背包总复杂度 O(NK2)O(NK^2)，空间复杂度 O(NK)O(NK)。

总结

这题的核心转化是：

  • 不是去直接算点对距离
  • 而是把点对贡献拆到每一条边上

一旦看出这一步，后面就是很标准的树上背包。

一图流解析

这张图把本题的建模、关键转移、实现检查和训练方法压缩到一页，适合读完正文后复盘。

一图流解析