习题： Paint the Tree（树DP）

题目

传送门

思路

因为颜色是无限的，我们可以轻松的转换一下题目

对于一个点，我们选取与他相连的k个点，并且将边权作为贡献

有了这个之后，我们考虑单纯的考虑一个子树是不行的，毕竟根节点与父亲也可能有关系

所以我们将这种情况加入\(dp\)状态的设计中来

设\(dp[i][0/1]\)表示以i为根节点的子树，i和\(k\)个点进行匹配还是\(k-1\)个点进行匹配的最大贡献

之后我们贪心地选取即可，这里的排序是用差值进行排序

设\(w_i\)为\(i\)和\(fa_i\)的边权的值

差值即为\(dp[v][1]+w_v-dp[v][0]\)

代码

#include<iostream>
#include<algorithm>
#include<vector>
#include<queue>
using namespace std;
struct node
{
    int e;
    int w;
    friend bool operator < (const node &a,const node &b)
    {
        return a.w>b.w;
    }
};
int t;
int n,k;
long long dp[500005][2];
vector<node> g[500005];
/*
以i为根节点的子树
0/1:0表示在子树内用了k种颜色，1表示在子树内用了k-1中颜色
*/
void dfs(int u,int fa)
{
    priority_queue<long long> q;
    if(g[u].size()==1&&g[u].front().e==fa)
        return;
    for(int i=0;i<g[u].size();i++)
    {
        int v=g[u][i].e;
        if(v!=fa)
        {
            dfs(v,u);
            dp[u][0]+=dp[v][0];
            dp[u][1]+=dp[v][0];
            q.push(dp[v][1]+g[u][i].w-dp[v][0]);
        }
    }
    for(int i=1;i<=k&&!q.empty();i++)
    {
        if(q.top()<0)
            break;
        if(i!=k)
            dp[u][1]=dp[u][1]+q.top();
        dp[u][0]=dp[u][0]+q.top();
        q.pop();
    }
}
void c_in()
{
    cin>>n>>k;
    for(int i=1;i<=n;i++)
    {
        dp[i][0]=dp[i][1]=0;
        g[i].clear();
    }
    for(int i=1,u,v,w;i<n;i++)
    {
        cin>>u>>v>>w;
        g[u].push_back((node){v,w});
        g[v].push_back((node){u,w});
    }   
    dfs(1,0);
    cout<<dp[1][0]<<endl;
}
int main()
{
    ios::sync_with_stdio(false);
    cin>>t;
    for(int i=1;i<=t;i++)
        c_in();
    return 0;
}