bzoj 4033: [HAOI2015]树上染色 树形dp

题意

有一棵点数为N的树，树边有边权。给你一个在0~N之内的正整数K，你要在这棵树中选择K个点，将其染成黑色，并将其他的N-K个点染成白色。将所有点染色后，你会获得黑点两两之间的距离加上白点两两之间距离的和的收益。问收益最大值是多少。
N<=2000,0<=K<=N

分析

这题太强啦。。。
先附一波题解
就是我们设f[i,j]表示以i为根的子树内选了j个黑点后子树内两两同色点对的距离和的话显然没法转移，因为没法考虑到子树外的状况。
那么我们就考虑换一种定义方式。
首先显然的是一条边的贡献就是$len*(左边白点数*右边白点数+左边黑点数*右边黑点数)$
那么我们就考虑一条边的总贡献。
设f[i,j]表示以i为根的子树内选了j个黑点，这棵子树内所有边在全局下的总贡献是多少。
那么我们可以枚举其一个子节点，若在其中放k个黑点，那么连接子节点和该节点的边对答案的贡献就是
$len*(k*(m-k)+(size[k]-k)*(n-size[k]-m+k))$
至于为什么的话自己根据上面那条算贡献的式子yy一下即可。

代码

#include<iostream>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<algorithm>
using namespace std;

typedef long long LL;

const int N=2005;
const int inf=0x3f3f3f3f;

int n,m,cnt,last[N],size[N];
struct edge{int to,next,len;}e[N*2];
LL f[N][N],tmp[N];

void addedge(int u,int v,int len)
{
    e[++cnt].to=v;e[cnt].len=len;e[cnt].next=last[u];last[u]=cnt;
    e[++cnt].to=u;e[cnt].len=len;e[cnt].next=last[v];last[v]=cnt;
}

void dp(int x,int fa)
{
    size[x]=1;
    for (int i=last[x];i;i=e[i].next)
    {
        if (e[i].to==fa) continue;
        dp(e[i].to,x);
        fill(tmp,tmp+size[x]+size[e[i].to]+1,-inf);
        for (int j=0;j<=size[x];j++)
            for (int k=0;k<=size[e[i].to];k++)
                tmp[j+k]=max(tmp[j+k],f[x][j]+f[e[i].to][k]+(LL)e[i].len*((LL)k*(m-k)+(LL)(n-size[e[i].to]-m+k)*(size[e[i].to]-k)));
        size[x]+=size[e[i].to];
        for (int j=0;j<=size[x];j++) f[x][j]=tmp[j];
    }
}

int main()
{
    scanf("%d%d",&n,&m);
    for (int i=1;i<n;i++)
    {
        int x,y,len;
        scanf("%d%d%d",&x,&y,&len);
        addedge(x,y,len);
    }
    dp(1,0);
    printf("%lld",f[1][m]);
    return 0;
}