树DP 树上染色

[HAOI2015]树上染色
时间限制：1 s 内存限制：256 MB
【题目描述】
有一棵点数为N的树，树边有边权。给你一个在0~N之内的正整数K，你要在这棵树中选择K个点，将其染成黑色，并将其他的N-K个点染成白色。将所有点染色后，你会获得黑点两两之间的距离加上白点两两之间距离的和的收益。问收益最大值是多少。
【输入格式】
第一行两个整数N,K。
接下来N-1行每行三个正整数fr,to,dis，表示该树中存在一条长度为dis的边(fr,to)。输入保证所有点之间是联通的。
【输出格式】
输出一个正整数，表示收益的最大值。
【输入样例1】
3 1
1 2 1
1 3 2
【输出样例1】
3
【数据范围】

对于30%的数据，N<=20
对于50%的数据，N<=100
对于100%的数据，N<=2000,0<=K<=N

明显是树DP,但如果去考虑哪个点是黑点就会很不方便，而任意两点间的距离过的边是确定的，那么可以考虑处理边，考虑边对答案的贡献。对于某条边，就是l=（他左侧黑点×他右侧黑点+他左侧白点×他右侧白点）×边权。
那么考虑转移，我们来挨个向答案中添加子树，枚举当前子树中有i个黑点，已枚举完的子树中共有j个黑点。
g[i+j]=max f[x][j]+f[son][i]+边权×l;
之前的所有子树中选j个+当前子树中选i个+贡献
必须单开一个数组g存当前的，因为还要用的之前的f[x][]，所以不能马上转移。
看似是N^3，实际上枚举不用到M,只要到子树的大小即可。所以是N^2

#include<iostream>
#include<cstdlib>
#include<cstdio>
#include<cstring>
#include<algorithm>
#define ll long long
using namespace std;
int n,m,e,adj[2005],size[2005];
ll f[2005][2005],g[2005]; 
struct road{int v,next,l;}lu[2005*2];
inline void add(int u,int v,int l){lu[++e]=(road){v,adj[u],l};adj[u]=e;}
inline void dp(int x,int fa)
{
    size[x]=1;
    for(int i=adj[x];i;i=lu[i].next)
    {
        int to=lu[i].v;
        if(to==fa)continue;
        dp(to,x);
        int x1=min(size[x],m),x2=min(size[to],m);
        for(int j=0;j<=x1;j++)
            for(int k=0;k<=x2;k++)
                if(j+k<=m)
                {
                    ll l=(m-k)*k+(size[to]-k)*(n-m-(size[to]-k));
                    g[k+j]=max(g[k+j],f[x][j]+f[to][k]+l*lu[i].l);
                }
        for(int j=0;j<=m;j++)f[x][j]=g[j],g[j]=0;
        size[x]+=size[to];
    }
}
int main()
{
    freopen("haoi2015_t1.in","r",stdin);
    freopen("haoi2015_t1.out","w",stdout);
    cin>>n>>m;int x,y,z;
    for(int i=1;i<n;i++)
    {
        scanf("%d%d%d",&x,&y,&z);
        add(x,y,z);add(y,x,z);
    }
    dp(1,0);
    cout<<f[1][m];
}