BZOJ4033[HAOI2015] 树上染色 解题报告【树上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染黑就能获得最大收益。
解题报告
我们考虑对于一条边ed[i]，他的权值会对答案有什么影响？
我们令u=ed[i].u，v=ed[i].v，w=ed[i].w；
那么他的贡献就是 w*( size[u]中的黑点 * size[v]中的黑点 +size[u]中的白点 * size[v]中的白点)（）。
我们令dp[u][j]表示以u为根节点的子树染j个黑点的最优答案，最终答案就在dp[1][k]中。
代码如下：

#include<cstdio>#include<cstring>#include<algorithm>using namespace std;const int N=2000;struct edge{    int v,next;    long long w;}ed[2*N+5];int n,K;int head[N+5],num;long long dp[N+5][N+5],size[N+5],temp[N+5];void build(int u,int v,long long w){    ed[++num].v=v;    ed[num].w=w;    ed[num].next=head[u];    head[u]=num;}void dfs(int u,int f){    size[u]=1;    for(int i=head[u];i!=-1;i=ed[i].next)    {        int v=ed[i].v;        if(v==f)continue;        dfs(v,u);        memset(temp,0,sizeof(temp));        for(int j=0;j<=size[u];j++)        for(int k=0;k<=size[v];k++)        temp[j+k]=max(temp[j+k],dp[u][j]+dp[v][k]+ed[i].w*(k*(K-k)+(size[v]-k)*(n-K-(size[v]-k))));        size[u]+=size[v];        for(int i=0;i<=size[u];i++)dp[u][i]=temp[i];    }}int main(){    scanf("%d%d",&n,&K);    memset(head,-1,sizeof(head));    for(int i=1;i<=n-1;i++)    {        int u,v,w;        scanf("%d%d%d",&u,&v,&w);        build(u,v,w);        build(v,u,w);    }    dfs(1,0);    printf("%lld\n",dp[1][K]);    return 0;}