Tree

1007 Tree

这道题的解题思路就是树形 dp，在求解过程中，尤其要注意 k 为0和1的时候的情况讨论。细节决定了这道题能不能A。

// Created by CAD
#include <bits/stdc++.h>
#define ll long long
using namespace std;

const int maxn=2e5+5;
struct pii{
    int fi;
    ll se,m;
};
vector<pii> g[maxn];
struct state{
    ll all=0,sumk=0;
    int siz=0;
    multiset<ll,greater<ll>> c;
    ll ck,ck_1;
}dp[maxn];
int k;
void dfs(int x,int fa,ll w){
    ll sum=0;
    auto &c=dp[x].c;
    for(pii &i:g[x])
        if(i.fi!=fa){
            dfs(i.fi,x,i.se);
            c.insert(dp[i.fi].sumk+i.se);
            i.m=dp[i.fi].sumk+i.se;
        }
    int cnt=0;
    dp[x].siz=c.size();
    for(ll i:c){
        if(++cnt<=k-1) dp[x].sumk+=i;
        if(cnt==k) dp[x].ck=i;
        if(cnt==k-1) dp[x].ck_1=i;
        dp[x].all+=i;
    }
}
ll ans=0;
void dfs(int x,int fa){
    ans=max(ans,dp[x].all);
    for(pii &i:g[x])
        if(i.fi!=fa){
            ll temp=dp[x].sumk,a=i.m;
            if(dp[x].siz<=k-1) temp+=-a+i.se;
            else if(dp[x].ck>=a) temp+=i.se;
            else temp+=-a+dp[x].ck+i.se;
            dp[i.fi].all+=temp;
            dp[i.fi].siz++;
            if(k-1==0);
            else if(dp[i.fi].siz<=k-1) dp[i.fi].sumk+=temp;
            else if(dp[i.fi].ck_1>=temp) dp[i.fi].ck=max(temp,dp[i.fi].ck);
            else
                dp[i.fi].sumk+=temp-dp[i.fi].ck_1,dp[i.fi].ck=dp[i.fi].ck_1;
            dfs(i.fi,x);
        }
}

int main() {
    int T;scanf("%d",&T);
    while(T--){
        int n;scanf("%d%d",&n,&k);
        for(int i=1;i<=n;++i)
            dp[i].sumk=dp[i].all=dp[i].ck=dp[i].ck_1=0,dp[i].c.clear(),g[i].clear();
        for(int i=1;i<=n-1;++i){
            int u,v;
            ll w;scanf("%d%d%lld",&u,&v,&w);
            g[u].push_back({v,w,0});
            g[v].push_back({u,w,0});
        }
        if(k==0){
            printf("0\n");
            continue;
        }
        ans=0;
        dfs(1,0,0);
        dfs(1,0);
        printf("%lld\n",ans);
    }
    return 0;
}