bzoj 4033: [HAOI2015]树上染色（树形DP）

4033: [HAOI2015]树上染色

Time Limit: 10 Sec   Memory Limit: 256 MB
Submit: 1786   Solved: 754
[ Submit][ Status][ Discuss]

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

dp[i][j]表示节点i的子树中染了k个点的最大收益

它的值并不是子树中所有相同颜色的点对距离之和，而是子树中每条边对答案的贡献之和

比如边(u, v) (u是v的儿子)对答案的贡献

==u子树中的黑点数*u子树外的黑点数*边长+u子树中的白点数*u子树外的白点数*边长

最后答案就是dp[1][k]，那么如何转移？

对于当前节点u求dp[u][d]，从左到右依次暴力所有的儿子，对于每个儿子暴力所有状态dp[v][c]

若c<d有dp[u][d] = max(dp[u][d], dp[v][c]+dp[u][d-c]+len(u, v)*c*(k-c)+len(u, v)*(siz[v]-c)*(n-k-siz[v]+c))

其中len(u, v)*c*(k-c)表示：子树v中的黑点数*u子树v外的黑点数*边长

len(u, v)*(siz[v]-c)*(n-k-siz[v]+c)表示：子树v中的白点数*u子树v外的白点数*边长

dp[u][d-c]是在遍历儿子v之前的其他儿子后求出来的

#include<stdio.h>
#include<string.h>
#include<vector>
#include<algorithm>
using namespace std;
#define LL long long
typedef struct
{
	LL v;
	LL len;
}Road;
vector<Road> G[2005];
Road now;
LL n, k, dp[2005][2005], siz[2005];
void Sech(LL u, LL p)
{
	LL i, v, c, d, len, full;
	siz[u] = 0;
	for(i=0;i<G[u].size();i++)
	{
		v = G[u][i].v;
		if(v!=p)
		{
			Sech(v, u);
			siz[u] += siz[v]+1;
		}
	}
	full = 1;
	for(i=0;i<G[u].size();i++)
	{
		v = G[u][i].v;
		len = G[u][i].len;
		if(v==p)
			continue;
		for(c=full;c>=0;c--)
		{
			for(d=min(k, siz[v]+1);d>=0;d--)
			{
				if(dp[v][d]>=0 && d+c<=k)
					dp[u][c+d] = max(dp[u][c+d], dp[v][d]+dp[u][c]+len*d*(k-d)+len*(siz[v]+1-d)*(n-k-siz[v]-1+d));
			}
		}
		full = min(k, full+siz[v]+1);
	}
}

int main(void)
{
	LL i, x, y;
	while(scanf("%lld%lld", &n, &k)!=EOF)
	{
		memset(dp, 0, sizeof(dp));
		for(i=1;i<=n;i++)
			G[i].clear();
		for(i=1;i<=n-1;i++)
		{
			scanf("%lld%lld%lld", &x, &y, &now.len);
			now.v = y;
			G[x].push_back(now);
			now.v = x;
			G[y].push_back(now);
		}
		Sech(1, 0);
		printf("%lld\n", dp[1][k]);
	}
	return 0;
}
/*
3 3
1 2 5
1 3 4
*/