BZOJ4033 树上染色

Problem 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

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

Example Input

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

Example Output

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

题解

考虑使用树形DP来做

• DP状态表示为dp[x][j],表示在以x为根的子树中选取j个点作为黑点的最大值，siz[x]表示以x为根的子树的大小。
• 枚举子节点y得到一条父亲节点与子节点的一条边，不妨假设边的权值为c
• 枚举子节点上选取黑点的个数为l,则这条边产生的贡献为两边黑点的个数相乘加上两边白点的个数相乘再乘上边的权值，即 $c*(l*(k-l)+(siz[y]-l)*(n-k-(siz[y]-l)))$
• 枚举黑点个数时应从大到小枚举避免重复计算

C++

#include<iostream>
#include<algorithm>
#include<cstdio>
#include<cmath>
#include<cstring>
#include<queue>
#include<set>
#include<map>
#include<stack>

using namespace std;

const long long N=2e3+50;

long long first[N],nxt[N*2],to[N*2],c[N*2],len,n,m,siz[N],dp[N][N],k;

void add(long long x,long long y,long long z)
{
    nxt[++len]=first[x];
    first[x]=len;
    to[len]=y;
    c[len]=z;
}

void dfs(int x,int fa)
{
    siz[x]=1;
    for (int i=first[x];i;i=nxt[i])
    if (to[i]!=fa)
    {
        int y=to[i];
        dfs(y,x);

        for (int j=min(siz[x],k);j>=0;j--)
        {
            for (int l=min(siz[y],k);l>=0;l--)
            {
                dp[x][j+l]=max(dp[x][j+l],dp[y][l]+dp[x][j]+c[i]*(l*(k-l)+(siz[y]-l)*(n-k-(siz[y]-l))));
            }
        }
        siz[x]+=siz[y];
    }
}

signed main()
{
#ifdef ONLINE_JUDGE
#else
    freopen("r.txt", "r", stdin);
#endif

    cin>>n>>k;
    for (int a,b,c,i=1;i<n;i++)
    {
        cin>>a>>b>>c;
        add(a,b,c);
        add(b,a,c);
    }

    dfs(1,0);

    cout<<dp[1][k]<<endl;
}