Helvetic Coding Contest 2017 online mirror K. Send the Fool Further! (medium)(树形DP)

题目链接：Helvetic Coding Contest 2017 online mirror K. Send the Fool Further! (medium)

题意：

给你一棵树,每条边有一个价值，现在每个节点最多访问k次，问最大的价值是多少。

每个价值只能加一次。

题解：

考虑dp[i][j(0||1)]表示以第i个节点为根的子树，最多访问i节点k次（0表示最后回到i节点，1表示不回到i节点）的最大价值。

状态转移方程为dp[i][0]=max{k-1个子节点的sum(dp[j][0])},dp[i][1]=max{k-1个子节点的sum(dp[j][0])+dp[l][1]}

#include<bits/stdc++.h>
#define F(i,a,b) for(int i=(a);i<=(b);++i)
using namespace std;
typedef pair<int,int>P;

const int N=1e5+7;
int n,k,ed;
P dp[N],tmp[N];
vector<P>g[N];

inline void up(int &a,int b){if(a<b)a=b;}

bool cmp(const P &a,const P &b){return a.first!=b.first?a.first>b.first:a.second>b.second;}

void dfs(int x=0,int fa=-1)
{
    for(auto &it:g[x])if(it.first!=fa)dfs(it.first,x);
    ed=0;
    for(auto &it:g[x])if(it.first!=fa)tmp[++ed]=P(dp[it.first].first+it.second,dp[it.first].second+it.second);
    if(!ed)return;
    sort(tmp+1,tmp+1+ed,cmp);
    for(int i=1;i<k&&i<=ed;i++)dp[x].first+=tmp[i].first;
    int sum=0;
    for(int i=1;i<=k&&i<=ed;i++)sum+=tmp[i].first;
    F(i,1,ed)
    {
        if(i<=k)up(dp[x].second,sum+tmp[i].second-tmp[i].first);
        else up(dp[x].second,sum+tmp[i].second-tmp[k].first);
    }
}

int main(){
    scanf("%d%d",&n,&k);
    F(i,1,n-1)
    {
        int x,y,z;
        scanf("%d%d%d",&x,&y,&z);
        g[x].push_back({y,z});
        g[y].push_back({x,z});
    }
    dfs();
    printf("%d\n",max(dp[0].second,dp[0].first));
    return 0;
}