BZOJ4033: [HAOI2015]树上染色（树形DP）

4033: [HAOI2015]树上染色

Time Limit: 10 Sec  Memory Limit: 256 MB
Submit: 3461  Solved: 1473
[Submit][Status][Discuss]

Description

有一棵点数为N的树，树边有边权。给你一个在0~N之内的正整数K，你要在这棵树中选择K个点，将其染成黑色，并

将其他的N-K个点染成白色。将所有点染色后，你会获得黑点两两之间的距离加上白点两两之间距离的和的收益。

问收益最大值是多少。

Input

第一行两个整数N,K。

接下来N-1行每行三个正整数fr,to,dis，表示该树中存在一条长度为dis的边(fr,to)。

输入保证所有点之间是联通的。

N<=2000,0<=K<=N

Output

输出一个正整数，表示收益的最大值。

Sample Input

5 2
1 2 3
1 5 1
2 3 1
2 4 2

Sample Output

17
【样例解释】
将点1,2染黑就能获得最大收益。

HINT

2017.9.12新加数据一组 By GXZlegend

Source

鸣谢bhiaibogf提供

思路：我们知道有道题，是求两两距离之和，我们可以一遍DFS得到，即每条边的贡献=sz[v]*(N-sz[v])*Len[i]；此题也一样，我们用dp[u][x]表示以u为根的子树里有x个黑点的最大收益。  不难得到DP方程。

注意这样的背包复杂度会偏大：

 for(int i=min(sz[u],K);i>=0;i--){
            for(int j=0;j<=min(i,sz[v]);j++){
                dp[u][i]=max(dp[u][i],dp[u][i-j]+dp[v][j]+1ll*j*(K-j)*Len[p[e]]+1ll*(sz[v]-j)*(N-K-sz[v]+j)*Len[p[e]]);
            }
        }

应该用如下：788ms

#include<bits/stdc++.h>
#define rep(i,a,b) for(int i=a;i<=b;i++)
#define ll long long
using namespace std;
const int maxn=2010;
int Laxt[maxn],Next[maxn<<1],To[maxn<<1],Len[maxn<<1];
int N,K,sz[maxn],cnt; ll dp[maxn][maxn];
void add(int u,int v,int w){
    Next[++cnt]=Laxt[u]; Laxt[u]=cnt; To[cnt]=v; Len[cnt]=w;
}
void dfs(int u,int f){
    sz[u]=1; dp[u][0]=dp[u][1]=0;
    for(int e=Laxt[u];e;e=Next[e]){
        int v=To[e]; if(v==f) return ;
        dfs(v,u);
        for(int i=min(sz[u],K);i>=0;i--){
            for(int j=min(sz[v],K-i);j>=0;j--){
                dp[u][i+j]=max(dp[u][i+j],dp[u][i]+dp[v][j]+1ll*j*(K-j)*Len[e]+1ll*(sz[v]-j)*(N-K-sz[v]+j)*Len[e]);
            }
        }
        sz[u]+=sz[v];
    }
}
int main()
{
    int u,v,w;
    scanf("%d%d",&N,&K);
    rep(i,1,N-1){
        scanf("%d%d%d",&u,&v,&w);
        add(u,v,w); add(v,u,w);
    }
    dfs(1,0);
    printf("%lld
",dp[1][K]);
    return 0;
}

也可以先排序：692ms

#include<bits/stdc++.h>
#define rep(i,a,b) for(int i=a;i<=b;i++)
#define ll long long
using namespace std;
const int maxn=2010;
int Laxt[maxn],Next[maxn<<1],To[maxn<<1],Len[maxn<<1];
int N,K,sz[maxn],cnt,son[maxn]; ll dp[maxn][maxn];
bool cmp(int w,int v) {return sz[To[w]]<sz[To[v]]; }
void add(int u,int v,int w){
    Next[++cnt]=Laxt[u]; Laxt[u]=cnt; To[cnt]=v; Len[cnt]=w;
}
void dfs1(int u,int f){
    sz[u]=1; for(int e=Laxt[u];e;e=Next[e]){
        if(To[e]!=f) dfs1(To[e],u),sz[u]+=sz[To[e]];
    }
}
void dfs2(int u,int f){
    int tot=0,p[maxn]; dp[u][0]=dp[u][1]=0; son[u]=1;
    for(int e=Laxt[u];e;e=Next[e]) if(To[e]!=f) p[++tot]=e;
    sort(p+1,p+tot+1,cmp);
    for(int e=1;e<=tot;e++){
        int v=To[p[e]]; if(v==f) return ;
        dfs2(v,u);
        for(int i=min(son[u],K);i>=0;i--){
            for(int j=min(son[v],K-i);j>=0;j--){
                dp[u][i+j]=max(dp[u][i+j],dp[u][i]+dp[v][j]+1ll*j*(K-j)*Len[p[e]]+1ll*(sz[v]-j)*(N-K-sz[v]+j)*Len[p[e]]);
            }
        }
        son[u]+=son[v];
    }
}
int main()
{
    int u,v,w;
    scanf("%d%d",&N,&K);
    rep(i,1,N-1){
        scanf("%d%d%d",&u,&v,&w);
        add(u,v,w); add(v,u,w);
    }
    dfs1(1,0); dfs2(1,0);
    printf("%lld
",dp[1][K]);
    return 0;
}