BZOJ 4033 [HAOI2015]树上染色

题目链接：传送门
$4033: [HAOI2015]树上染色$

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

【样例解释】

将点1,2染黑就能获得最大收益。

这道题很神奇啊
让着把k个点染成黑色
求每个黑点与黑点和白点与白点之间的距离和
转化一下
我们把距离转成边
我们求出每条边被经过的次数
最后乘上每条边的权值就好了
下面一步转化十分重要
就是求每条边被经过的次数
考虑两个同色的点
如果一条边不在这两个点之间的路径上
我们不用管它
如果在的话
它就会对答案产生贡献了
也就是如果两个同色点在一条边的两侧
那这个边就会被经过一次
最后得出
一条边被经过的次数就等于边的两侧同色点个数的乘积
于是下面的式子就有了

$cnt=k∗(kk−k)+(siz[ca]−k)∗(n−kk−siz[ca]+k)$
$k*(kk-k)$ 就是一边的黑点个数乘以另一边的黑点个数
$(siz[ca]-k)$ 就是一边的白点个数，后面的就是另一边的白点个数

$siz$ 是当前节点的子树大小， $kk$ 是要选的黑点数， $k$ 是当前子节点的子树上已经选的黑点数， $ca$ 是遍历到的节点，也就是 $edge[i].to$
转移方程也很显然

$f[fr][j]=max(f[fr][j],f[fr][j−k]+f[ca][k]+cnt∗edge[i].dis)$

$fr$ 为当前节点， $j$ 为枚举的黑点个数， $cnt$ 就是一条边被经过的次数
注意如果要选的黑色点超过了节点数的一半，要把黑色节点数换成另一块
不然会 $T$ 很惨（ $Luogu$ 除外）

#include <iostream>
#include <cstdio>
#include <cstring>
#include <cstdlib>
#include <complex>
#include <algorithm>
#include <climits>
#include <queue>
#include <map>
#include <vector>
#include <iomanip>
#define A 1000010
#define B 2010
#define ll long long

using namespace std;
struct node {
	int next, to;
	ll dis;
}edge[A];
int head[A], num_edge;
void add_edge(int from, int to, ll dis) {
	edge[++num_edge].next = head[from];
	edge[num_edge].to = to;
	edge[num_edge].dis = dis;
	head[from] = num_edge;
}
int n, kk, a, siz[A], b;
ll f[B][B], c;
void dfs(int fr, int fa) {
	siz[fr] = 1;
	for (int i = head[fr]; i; i = edge[i].next) {
		int ca = edge[i].to;
		if (ca == fa) continue;
		dfs(ca, fr);
		siz[fr] += siz[ca];
		for (int j = min(siz[fr], kk); j >= 0; j--) {
			if (f[fr][j] != -1) f[fr][j] += f[ca][0] + siz[ca] * (n - kk - siz[ca]) * edge[i].dis;
			for (int k = min(j, siz[ca]); k > 0; k--)
				if (f[fr][j - k] != -1)
			  	f[fr][j] = max(f[fr][j], f[fr][j - k] + f[ca][k] + (k * (kk - k) + (siz[ca] - k) * (n - kk - siz[ca] + k)) * edge[i].dis);
		}
	}
}

int main() {
	cin >> n >> kk;
	if (n - kk < kk) kk = n - kk;
	for (int i = 1; i < n; i++) {
		cin >> a >> b >> c;
		add_edge(a, b, c);
		add_edge(b, a, c);
	}
	memset(f, -1, sizeof f);
	for (int i = 1; i <= n; i++) f[i][0] = f[i][1] = 0;
	dfs(1, 0);
	printf("%lld\n", f[1][kk]);
}