bzoj4033 树上染色

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

 

其实蛮正常的一个树dp，然而我狂T。然后发现一个for循环中k可以写成min( k , size[pos] )稍微优化一下，然后就过了呵呵。

//Serene
#include<algorithm>
#include<iostream>
#include<cstring>
#include<cstdlib>
#include<cstdio>
#include<cmath>
using namespace std;
const int maxn=2000+10,INF=2e8;
int n,k;
long long dp[maxn][maxn],ans;

int aa;char cc;
int read() {
	aa=0;cc=getchar();
	while(cc<'0'||cc>'9') cc=getchar();
	while(cc>='0'&&cc<='9') aa=aa*10+cc-'0',cc=getchar();
	return aa;
}

int fir[maxn],nxt[2*maxn],to[2*maxn],v[2*maxn],e=0;
void add(int x,int y,int z) {
	to[++e]=y;nxt[e]=fir[x];fir[x]=e;v[e]=z;
	to[++e]=x;nxt[e]=fir[y];fir[y]=e;v[e]=z;
}

int fa[maxn],size[maxn];
void dfs(int pos,int dis) {
	size[pos]=1;dp[pos][0]=dp[pos][1]=0;
	for(int y=fir[pos];y;y=nxt[y]) {
		if(to[y]==fa[pos]) continue;
		fa[to[y]]=pos;
		dfs(to[y],v[y]);
		size[pos]+=size[to[y]];
		for(int i=min(k,size[pos]);i>=0;--i)
			for(int j=0;j<=i&&j<=size[to[y]];++j) {
				dp[pos][i]=max(dp[pos][i],dp[pos][i-j]+dp[to[y]][j]);
		}
	}
	if(dis) for(int i=0;i<=k;++i) dp[pos][i]+=(long long)dis*((long long)i*(k-i)+(long long)(size[pos]-i)*(n-k-size[pos]+i));
}

int main() {
	n=read();k=read(); int x,y,z;
	for(int i=1;i<n;++i) {
		x=read();y=read();z=read();
		add(x,y,z);
	}
	for(int i=1;i<=k;++i) for(int j=1;j<=n;++j) dp[j][i]=0-(long long)INF;
	dfs(1,0);
	printf("%lld",dp[1][k]);
	return 0;
}