CF 1223E Paint the Tree

题目描述

https://codeforces.com/problemset/problem/1223/E

简要题意：给定一棵有 $n$ 个点的无根树和 $k$​，每个点要染上恰好 $k$ 种颜色，现在有无数种颜色，每种颜色最多用两次，当一条边的两个端点附上的颜色中有至少一种相同颜色时，这条边的贡献就是这条边的权值，否则是 $0$，求一种染色方案，使得这棵树的所有边的贡献之和最大

Solution

题目可以转化为在树上选取若干条边，需要保证加入这些边后没有一个点的度超过 $k$，问最大边权和

对于儿子，我们只需要知道儿子还能不能选一次，所以我们令 $f[u]$ 表示还能选，$g[u]$ 表示不能选

在转移的时候，我们首先都选 $g[v]$，然后将 $f[v]+w-g[v]$ 排序选前 $k$ 个作为 $g[u]$，前 $k-1$ 个作为 $f[u]$

#include <iostream>
#include <vector>
#include <algorithm>
#define maxn 500010
#define ll long long
using namespace std;

int n, k;

struct Edge {
    int to, next, w; 
} e[maxn * 2]; int c1, head[maxn];
inline void add_edge(int u, int v, int w) {
    e[c1].to = v; e[c1].w = w; 
    e[c1].next = head[u]; head[u] = c1++;
}

ll f[maxn], g[maxn];
void dfs(int u, int fa) {
    g[u] = 0; int t = k; vector<ll> A; 
    for (int i = head[u]; ~i; i = e[i].next) {
        int v = e[i].to, w = e[i].w; if (v == fa) continue;
        dfs(v, u); g[u] += g[v]; A.push_back(f[v] + w - g[v]); 
    }
    sort(A.begin(), A.end(), greater<ll>());
    f[u] = g[u]; 
    for (auto v : A) {
        if (v < 0 || !t) break;
        g[u] += v; --t;
        if (t >= 1) f[u] = g[u];
    } 
} 

void work() { 
    cin >> n >> k;
    fill(head + 1, head + n + 1, -1); c1 = 0; 
    for (int i = 1; i < n; ++i) {
        int x, y, z; cin >> x >> y >> z;
        add_edge(x, y, z); add_edge(y, x, z); 
    } dfs(1, 0); cout << max(g[1], f[1]) << "\n"; 
}

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

    int T; cin >> T;
    while (T--) work(); 
    return 0;
}