HAOI2015」树上染色

有一棵点数为 N 的树，树边有边权。给你一个在 0∼N 之内的正整数K，你要在这棵树中选择 K 个点，将其染成黑色，并将其他的 N−K个点染成白色。将所有点染色后，你会获得黑点两两之间的距离加上白点两两之间距离的和的收益。

问收益最大值是多少。

输入
第一行两个整数 N,K。

接下来 N−1 行每行三个正整数 fr,to,dis，表示该树中存在一条长度为 dis 的边 (fr,to)

输入保证所有点之间是联通的。

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

样例输入 [复制]
5 2
1 2 3
1 5 1
2 3 1
2 4 2
样例输出 [复制]
17

思路
树形DP
size[i]:i子树的节点个数
f[i][j]:在i子树中染j个黑点的最大贡献
更新时考虑每条边对答案的贡献
即：这条边两侧的黑点个数乘积边权+两侧白点个数乘积边权

#include<bits/stdc++.h>
using namespace std;
#define int long long
inline int read()
{
    int data=0;int w=1; char ch=0;
    ch=getchar();
    while(ch!='-' && (ch<'0' || ch>'9')) ch=getchar();
    if(ch=='-') w=-1,ch=getchar();
    while(ch>='0' && ch<='9') data=(data<<3)+(data<<1)+ch-'0',ch=getchar();
    return data*w;
}
const int N=2010;
int n,k;
struct node{
	int u,v,w,nxt;
}e[N<<1];
int fir[N],cnt=0,dp[N][N],siz[N];
inline void add(int u,int v,int w){
	e[++cnt]=(node){u,v,w,fir[u]};fir[u]=cnt;
} 
inline int calc(int w,int siz,int y){
	  int ans=1;
      ans*=w*(y*(k-y)+(n-siz-k+y)*(siz-y));
      return ans;
}
inline void dfs(int u,int fa){
	siz[u]=1;
	for(int i=fir[u];i;i=e[i].nxt){
		int v=e[i].v;
		if(v==fa) continue;
		dfs(v,u);
		for(int x=siz[u];x>=0;x--){
			for(int y=siz[v];y>=0;y--){
				int sum=dp[u][x]+dp[v][y]+calc(e[i].w,siz[v],y);
				dp[u][x+y]=max(dp[u][x+y],sum);
			} 
		} 
		siz[u]+=siz[v];
	}
} 
signed main(){
    n=read();k=read();
    for(int i=1;i<n;i++){
    	int u=read(),v=read(),w=read();
    	add(u,v,w);add(v,u,w);
	}
	dfs(1,0);
	int ans=dp[1][k];
	printf("%lld",ans);
	return 0;
}