【BZOJ4033】【HAOI2015】树上染色

Description

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

Input

第一行包含两个整数 N, K 。接下来 N-1 行每行三个正整数 fr, to, dis ， 表示该树中存在一条长度为 dis 的边 (fr, to) 。输入保证所有点之间是联通的。

Output

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

Sample Input

3 1

1 2 1

1 3 2

Sample Output

3

Hint

对于 100% 的数据，\(0 \leq K \leq N \leq 2000\)

Solution

考虑树形DP，状态应为第i个点的子树内有j个黑色点对整棵树的贡献，问题就转换为了树形背包，对于一条边的贡献=两边的黑色点数之积+两边的白色点之积，利用这一点进行树形DP即可，每次枚举子树内黑色点个数，与当前子树内黑色点个数，进行统计答案即可，将上界改为子树大小，可以将时间复杂度优化至\(O(n^2)\).

Code

#include <stdio.h>
#include <string.h>
#define MN 2005
#define R register
#define ll long long
#define file(x) freopen(#x".in","r",stdin);freopen(#x".out","w",stdout);
#define end fclose(stdin);fclose(stdout)
inline int read(){
	R int x; R bool f; R char c;
	for (f=0; (c=getchar())<'0'||c>'9'; f=c=='-');
	for (x=c-'0'; (c=getchar())>='0'&&c<='9'; x=(x<<3)+(x<<1)+c-'0');
	return f?-x:x;
}
int n,k,to[MN<<1],nt[MN<<1],v[MN<<1],h[MN],sz[MN],en;ll f[MN][MN];
inline void ins(int x,int y,int vl){to[++en]=y,nt[en]=h[x],h[x]=en,v[en]=vl;}
inline int min(int a,int b){return a<b?a:b;}
inline ll max(ll a,ll b){return a>b?a:b;}
inline void dp(int u,int fa){
	f[u][0]=f[u][1]=0;sz[u]=1;
	for (R int i=h[u]; i; i=nt[i])
		if (to[i]!=fa) dp(to[i],u),sz[u]+=sz[to[i]];
	for (R int i=h[u]; i; i=nt[i])
		if (to[i]!=fa)
			for (R int j=min(k,sz[u]); j>=0; --j)
				for (R int l=0; l<=min(j,sz[to[i]]); ++l)
					if (~f[u][j-l]){
						ll val=(1ll*l*(k-l)+(1ll*sz[to[i]]-l)*(1ll*n-k+l-sz[to[i]]))*v[i];
						f[u][j]=max(f[u][j],f[u][j-l]+f[to[i]][l]+val);
					}
}
int main(){
	memset(f,-1,sizeof(f));
	n=read(),k=read();
	for (R int i=1; i<n; ++i){
		R int x=read(),y=read(),v=read();
		ins(x,y,v); ins(y,x,v);
	}dp(1,0);printf("%lld\n",f[1][k]);
	return 0;
}