2020杭电多校第五场 Tree（换根DP）

Problem Description
Aphelios likes to play with tree. A weighted tree is an undirected connected graph without cycles and each edge has a weight. The degree of each vertex is the number of vertexes which connect with it. Now Aphelios has a weighted tree T with n vertex and an integer k, and now he wants to find a subgraph G of the tree, which satisfies the following conditions:

1. G should be a connected graph, in other words, each vertex can reach any other vertex in the subgraph G.

2. the number of the vertex whose degree is larger than k is no more than 1.

3. the total weight of the subgraph is as large as possiblie.

Now output the maximum total weight of the subgraph you can find.

Input
The first line contains an integer q, which represents the number of test cases.

For each test case, the first line contains two integer n and k (1≤n≤2×105,0≤k<n).

For next n−1 lines , each line contains 3 numbers u,v,d, which means that there is an edge between u and v weighted d (1≤u,v≤n,0≤d≤109).

You may assume that ∑n≤106.

Output
For each test case, output one line containing a single integer ans denoting the total weight of the subgraph.

Sample Input
1
5 2
1 2 5
2 3 2
2 4 3
2 5 4

Sample Output
14

Source
2020 Multi-University Training Contest 5

题意：
树中选择一个子图，使得图中点度数大于K的至多只有一个。求这个子图边权和最大值。

思路：
裸的换根DP
定义$f[i][0]$为以$i$为根节点且所选子树度数和小于等于k-1的最大边权和。
$f[i][1]$为以$i$为根节点且所选子树度数和小于等于k-1的最大边权和。

那么$f[i][1]$转移就是取子树最大的前k个$f[v][0]+w$值加上，$f[i][0]$就是取最大的前k-1个，标记选了哪些点。第二次dfs的时候再考虑父节点过来的贡献，只是需要注意这个点是否被属于父亲的前k大节点。

两次dfs就可以把$f$数组给维护出来。
注意到题目的还可以允许一个点度数大于$k$，那我们就枚举所有点的$f[u][1]$，再加上这个点没有选的其他子节点的$f[v][0]+w$。结果取最大就好了。

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <cmath>
#include <vector>
#include <queue>
#include <iostream>
#include <map>
#include <string>

using namespace std;

typedef long long ll;

const int mod = 998244353;
const int maxn = 4e5 + 7;

ll f[maxn][2];
int head[maxn],nex[maxn],to[maxn],val[maxn],tot;
int vis[maxn]; //前k个的标记
int n,k;
ll ans;

void add(int x,int y,int z) {
    to[++tot] = y;
    nex[tot] = head[x];
    val[tot] = z;
    head[x] = tot;
}

void init() {
    for(int i = 1;i <= n;i++) {
        head[i] = 0;
        f[i][0] = f[i][1] = 0;
        vis[i] = 0;
    }
    tot = 0;
}

void DP1(int x,int fa) {
    priority_queue<pair<ll,int>>q;
    for(int i = head[x];i;i = nex[i]) {
        int v = to[i],w = val[i];
        if(v == fa) continue;
        DP1(v,x);
        q.push({-(f[v][0] + w),v});
        if(q.size() > k) {
            q.pop();
        }
    }
    if(q.size() == k) {
        f[x][1] += -q.top().first;
        vis[q.top().second] = 1;
        q.pop();
    }
    while(!q.empty()) {
        ll num = -q.top().first,id = q.top().second;
        f[x][1] += num;
        f[x][0] += num;
        vis[id] = 1;
        q.pop();
    }
}

void DP2(int x,int fa) {
    for(int i = head[x];i;i = nex[i]) {
        vis[to[i]] = 0;
    }
    priority_queue<pair<ll,int> >q;
    
    for(int i = head[x];i;i = nex[i]) {
        int v = to[i],w = val[i];
        if(v == fa) {
            if(vis[x]) { //x在父亲的前k个里面
                q.push({-(f[fa][1] - f[x][0]),v});
            } else {
                q.push({-(f[fa][0] + w),v});
            }
        } else {
            q.push({-(f[v][0] + w),v});
        }
        if(q.size() > k) q.pop();
    }
    
    ll sum1 = 0,sum2 = 0;
    
    if(q.size() == k) {
        sum1 += -q.top().first;
        vis[q.top().second] = 1;
        q.pop();
    }
    while(!q.empty()) {
        ll num = -q.top().first,id = q.top().second;
        sum1 += num;
        sum2 += num;
        vis[id] = 1;
        q.pop();
    }
    
    ll sum = 0;
    for(int i = head[x];i;i = nex[i]) {
        int v = to[i],w = val[i];
        if(!vis[v]) {
            if(v == fa) {
                if(vis[x]) {
                    sum += f[v][1] - f[x][0];
                } else {
                    sum += w + f[v][0];
                }
            } else {
                sum += w + f[v][0];
            }
        }
    }
    
    f[x][0] = max(f[x][0],sum2);
    f[x][1] = max(f[x][1],sum1);
    
    sum += f[x][1];
    
    ans = max(ans,sum);
    
    for(int i = head[x];i;i = nex[i]) {
        int v = to[i];
        if(v == fa) continue;
        DP2(v,x);
    }
}

int main() {
    int T;scanf("%d",&T);
    while(T--) {
        scanf("%d%d",&n,&k);
        init();
        for(int i = 1;i < n;i++) {
            int x,y,z;scanf("%d%d%d",&x,&y,&z);
            add(x,y,z);add(y,x,z);
        }
        if(k == 0) {
            printf("0\n");
            continue;
        }
        ans = 0;
        DP1(1,0);
        DP2(1,0);
        printf("%lld\n",ans);

    }
    return 0;
}