Codeforces Round #591 E. Paint the Tree 树形dp

昨晚写完这题，交了三次没交上…
E. Paint the Tree

题意：给你一颗带边权的树，每个节点使用次数均为$k$，你如果想要获得一条边的权值，那么必须要消耗该边相连的两个点的一次使用次数，问最多能获得多少的权值

解法：以1为根，设$d[u][0]，d[u][1]$分别为$u$子树中，$u$节点可使用次数为0和不为0所能获得的最大权值，假设$u$有$T$个儿子$son$，我们把$d[so{n}_i][0]，d[so{n}_i][1]+w_i，d[so{n}_i][1]$存到优先队列，设$W_i=d[so{n}_i][1]+w_i-d[so{n}_i][0]$，则优先队列按照$W_i$从大到小排序，如果有不超过$k-1$个的$W_i$大于0，我们最多只会消耗$u$节点$k-1$个使用次数，那么可以直接转移：$d[u][0]=d[u][1]={\sum }_{i=1}^{T}max(d[so{n}_i][0],d[so{n}_i[1]+w_i)$，如果有超过$k-1$个$W_i$大于0，那么$d[u][0]=({\sum }_{i=1}^{k}d[so{n}_i][1]+w_i)+{\sum }_{i=k+1}^{T}max(d[so{n}_i][0],d[so{n}_i][1])$，$d[u][1]=({\sum }_{i=1}^{k-1}d[so{n}_i][1]+w_i)+{\sum }_{i=k}^{T}max(d[so{n}_i][0],d[so{n}_i][1])$

#include<bits/stdc++.h>
#define pi pair<int, int>
#define mk make_pair
#define ll long long
using namespace std;
const int maxn = 5e5 + 10;
ll d[maxn][2];
int n, k;
vector<pi> G[maxn];
struct node {
    ll w1, w2, w3;
    bool operator<(const node& t) const {
        return w2 - w1 < t.w2 - t.w1;
    }
};
priority_queue<node> q;
void dfs(int u, int fa) {
    for (auto tmp : G[u]) {
        int v = tmp.first;
        int w = tmp.second;
        if (v == fa)
            continue;
        dfs(v, u);
    }
    for (auto tmp : G[u]) {
        int v = tmp.first;
        int w = tmp.second;
        if (v == fa)
            continue;
        q.push(node{d[v][0], d[v][1] + w, d[v][1]});
    }
    int flag = 0;
    for (int i = 1; i < k; i++) {
        if (q.empty())
            break;
        ll w = max(q.top().w2, q.top().w1);
        d[u][0] += w;
        d[u][1] += w;
        if (q.top().w1 == w)
            flag = 1;
        q.pop();
    }
    if (!q.empty()) {
        ll w = max(q.top().w2, q.top().w1);
        d[u][0] += w;
        d[u][1] += max(q.top().w1, q.top().w3);
        if (q.top().w1 == w)
            flag = 1;
        q.pop();
    }
    else
        flag = 1;
    while (!q.empty()) {
        d[u][0] += max(q.top().w1, q.top().w3);
        d[u][1] += max(q.top().w1, q.top().w3);
        q.pop();
    }
    if (flag)
        d[u][0] = d[u][1] = max(d[u][0], d[u][1]);
}
void solve() {
    int u, v, w;
    scanf("%d%d", &n, &k);
    for (int i = 1; i < n; i++) {
        scanf("%d%d%d", &u, &v, &w);
        G[u].push_back(mk(v, w));
        G[v].push_back(mk(u, w));
    }
    dfs(1, 0);
    printf("%lld\n", max(d[1][0], d[1][1]));
    for (int i = 1; i <= n; i++)
        G[i].clear(), d[i][0] = d[i][1] = 0;
}
int main() {
    int T;
    scanf("%d", &T);
    while (T--)
        solve();
}