【HAOI2015】树上染色 - 树形 DP

题目描述

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

输入输出格式

输入格式：

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

输出格式：

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

思路

设 \(f_{i,j}\) 为以 \(i\) 为根的子树选 \(j\) 个点染成黑色的贡献。考虑一条边的贡献，它被一条路径经过就会产生贡献，经过这个点的路径的数量显然是左边黑色节点个数和右边黑色节点个数的积。就可以得出一个节点 \(u\) 被 \((j-m) \times j + (size_u-j) \times (n-m-size_u+j)\)，其中 \(j\) 为已经选了的黑色节点个数，\(size\) 为子树大小。直接转移就行了

/************************************************
*Author        :  lrj124
*Created Time  :  2019.03.18.20:14
*Mail          :  1584634848@qq.com
*Problem       :  luogu3177
************************************************/
#include <bits/stdc++.h>
using namespace std;
const int maxn = 2000 + 10;
struct Edge {
	int to,val;
	Edge(int v,int w) : to(v),val(w) {}
};
vector<Edge> edge[maxn];
long long f[maxn][maxn];
int n,K,size[maxn];
inline void dp(int now,int fa) {
	size[now] = 1;
	f[now][0] = f[now][1] = 0;
	for (size_t i = 0;i < edge[now].size();i++)
		if (edge[now][i].to ^ fa) {
			dp(edge[now][i].to,now);
			size[now] += size[edge[now][i].to];
		}
	for (size_t i = 0;i < edge[now].size();i++) if (edge[now][i].to ^ fa)
		for (int j = min(K,size[now]);j >= 0;j--)
			for (int k = 0;k <= min(j,size[edge[now][i].to]);k++)
				if (f[now][j-k] >= 0) f[now][j] = max(f[now][j],f[now][j-k]+f[edge[now][i].to][k]+1ll*k*(K-k)*edge[now][i].val+1ll*(size[edge[now][i].to]-k)*(n-size[edge[now][i].to]-K+k)*edge[now][i].val);
}
int main() {
	scanf("%d%d",&n,&K);
	memset(f,128,sizeof(f));
	for (int i = 1,u,v,w;i < n;i++) {
		scanf("%d%d%d",&u,&v,&w);
		edge[u].push_back(Edge(v,w));
		edge[v].push_back(Edge(u,w));
	}
	dp(1,0);
	printf("%lld",f[1][K]);
	return 0;
}