BZOJ 4033: [HAOI2015]树上染色 (树形DP)

题意

n个点的树，选k个黑点，求黑点之间两两距离之和加上白点之间两两距离之和的最大值。

题解

感觉两两距离之和没法统计。但是我们可以对每条边统计左边和右边各自每种颜色的点数，乘起来就是这条边要被统计到的次数。那么就直接树形DP，统计子树内选了几个黑点。然后向上传的时候加上父亲边的贡献。

一个技巧就是如果dp的二维状态的第二维是子树节点个数，那么合并两棵子树的时候直接暴力枚举转移就行了。每两个点只会在lca处被统计，所以是n^2的。

但是我写的时候细节出了点问题，10s的题9.9s过了。。。然后修改了一下就600ms了。
具体见代码。

CODE

#pragma GCC optimize (2)
#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
const int MAXN = 2005;
int n, m, fir[MAXN], to[MAXN<<1], nxt[MAXN<<1], cnt, sz[MAXN];
LL f[MAXN][MAXN], w[MAXN<<1];
inline void link(int u, int v, int wt) {
	to[++cnt] = v; nxt[cnt] = fir[u]; fir[u] = cnt; w[cnt] = wt;
	to[++cnt] = u; nxt[cnt] = fir[v]; fir[v] = cnt; w[cnt] = wt;
}
void dfs(int u, int ff) {
	sz[u] = 1; f[u][0] = f[u][1] = 0;
	for(int i = fir[u], v; i; i = nxt[i])
		if((v=to[i]) != ff) {
			dfs(v, u); sz[u] += sz[v];
			for(int j = 0; j <= sz[v] && j <= m; ++j)
				f[v][j] += (1ll * j * (m-j) + 1ll * (sz[v]-j) * ((n-m)-(sz[v]-j))) * w[i];
			for(int j = min(sz[u], m); j >= 0; --j)
				for(int k = max(j-sz[u]+sz[v], 0); k <= sz[v] && k <= j; ++k)
					f[u][j] = max(f[u][j], f[u][j-k] + f[v][k]); //j-k <= sz[u]-sz[v]
				//之前我这里的k是从0开始枚举的。其实有问题。
				//正确的复杂度是基于分别枚举两棵子树的节点个数
				//但我这里sz[u]已经加上了sz[v]，所以时间复杂度不对。
				//比较简单的方法是存一个tmp数组来暂时记录答案，最后再赋值给f[u]
				//我比较懒，不想写tmp数组，就直接转移到f[u]，如上，
				//                                   必须保证j-k小于等于u之前的子树节点之和
				//而我前面已经加上了sz[v]，所以写出来就是j-k<=sz[u]-sz[v]，即k>=j-sz[u]+sz[v]
				//这样就是严格n^2
		}
}
int main () {
	scanf("%d%d", &n, &m);
	for(int i = 1, u, v, wt; i < n; ++i)
		scanf("%d%d%d", &u, &v, &wt), link(u, v, wt);
	memset(f, -0x3f, sizeof f);
	dfs(1, 0);
	printf("%lld\n", f[1][m]);
}