hdu_6820 Tree

树形dp裸题

两遍dfs

第一次求dp[i][0]表示选取k-1条边（除父节点外）所能取得的最大权值

第二次求dp[i][1]表示i点可以超过k条边其他点不行所能取得的最大权值，过程中更新父亲的dp[fa][0]

#include <cstdio>
#include <algorithm>
#include <vector>
using namespace std;
#define ll long long
struct sons{ll dp;int id;};
struct edge{int v;ll w;};
vector<edge> G[200010];
int p[200010],n,k,T;
ll dp[200010][2],ans;

void init(){
    for(int i=1;i<=n;i++){G[i].clear(); dp[i][0]=dp[i][1]=0;}
    ans=0;
}

bool cmp1(ll x,ll y){return x>y;}
bool cmp(sons x,sons y){return x.dp>y.dp;}

void dfs(ll u){
    ll son[G[u].size()+10]; int cnt=0;
    for(int i=0;i<G[u].size();i++)if(G[u][i].v!=p[u]){
        int v=G[u][i].v; ll w=G[u][i].w;
        p[v]=u;
        dfs(v);
        son[++cnt]=dp[v][0]+w;
    }
    if(!cnt)return;
    sort(son+1,son+cnt+1,cmp1);
    for(int i=1;i<=min(cnt,k-1);i++)dp[u][0]+=son[i];
}

void dfs2(ll u){
    sons son[G[u].size()+10];int cnt=0;
    for(int i=0;i<G[u].size();i++){
        int v=G[u][i].v; ll w=G[u][i].w;
        if(G[u][i].v!=p[u])son[++cnt].dp=dp[v][0]+w;
        else son[++cnt].dp=dp[v][0]-dp[u][0];
        son[cnt].id=v;
    }
	if(!cnt)return;
    sort(son+1,son+cnt+1,cmp);
    for(int i=1;i<=cnt;i++)dp[u][1]+=son[i].dp;
    dp[u][0]=0;
    for(int i=1;i<=min(cnt,k-1);i++)dp[u][0]+=son[i].dp;
    //prllf("%d:%d %d
",u,dp[u][0],dp[u][1]);
    ans=max(ans,dp[u][1]);
    for(int i=1;i<=cnt;i++)if(son[i].id!=p[u]){
    	if(i<k){if(k<=cnt)dp[u][0]+=son[k].dp;}
    	else dp[u][0]+=son[i].dp;
    	//prllf("%d:%d %d
",u,dp[u][0],dp[u][1]);
        dfs2(son[i].id);
		if(i<k){if(k<=cnt)dp[u][0]-=son[k].dp;}
		else dp[u][0]-=son[i].dp;
    }
}

int main(){
    scanf("%d",&T);
    while(T--){
        scanf("%d%d",&n,&k);
        for(int i=1;i<=n;i++)p[i]=i;
        for(int i=1;i<n;i++){
            int u,v;ll w;
            scanf("%d%d%lld",&u,&v,&w);
            G[u].push_back((edge){v,w});
            G[v].push_back((edge){u,w});
        }
        if(k==0){printf("0
"); init(); continue;}
        dfs(1);
        dfs2(1);
        //for(ll i=1;i<=n;i++)prllf("%d %d
",dp[i][0],dp[i][1]);
        printf("%lld
",ans);
        init();
    }
    return 0;
}