[bzoj4033][树形DP]树上染色

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

题解

树形DP吧，我一开始状态想复杂了。
设f[i][j]表示以第i个点为根的子树中，有j个节点为黑色的总贡献
转移的话
我们枚举一个点与一条边，这个点之前的子树状态已经处理完毕，现在要加入这条边的贡献
很容易想到，这条边的贡献即为左边子树黑色点数乘右边子树黑色点数乘边长度，白色也同理
那么瞎转移即可

#include<cstdio>
#include<cstring>
#include<cstdlib>
#include<algorithm>
#include<cmath>
using namespace std;
typedef long long LL;
LL f[2100][2100];//第i个点 子树中j个黑点 本身颜色    总贡献
struct node
{
    int x,y,c,next;
}a[11000];int len,last[2100];
void ins(int x,int y,int c)
{
    len++;
    a[len].x=x;a[len].y=y;a[len].c=c;
    a[len].next=last[x];last[x]=len;
}
int n,kk,tot[2100],fa[2100];
void treedp(int x)
{
    tot[x]=1;
    for(int k=last[x];k;k=a[k].next)
    {
        int y=a[k].y;
        if(y==fa[x])continue;
        fa[y]=x;
        treedp(y);
        for(int i=tot[x];i>=0;i--)
        {
            for(int j=1;j<=tot[y];j++)
                f[x][i+j]=max(f[x][i+j],f[y][j]+f[x][i]+(LL)a[k].c*j*(kk-j)+(LL)a[k].c*(n-kk-(tot[y]-j))*(tot[y]-j));
            f[x][i]+=f[y][0]+(LL)a[k].c*(n-kk-tot[y])*tot[y];
        }
        tot[x]+=tot[y];
    }
}
int main()
{
    scanf("%d%d",&n,&kk);
    len=0;memset(last,0,sizeof(last));
    for(int i=1;i<n;i++)
    {
        int x,y,c;
        scanf("%d%d%d",&x,&y,&c);
        ins(x,y,c);ins(y,x,c);
    }
    fa[1]=0;
    treedp(1);
    printf("%lld\n",f[1][kk]);
    return 0;
}