[BZOJ4033][HAOI2015]树上染色

[BZOJ4033][HAOI2015]树上染色

试题描述

有一棵点数为N的树，树边有边权。给你一个在0~N之内的正整数K，你要在这棵树中选择K个点，将其染成黑色，并

将其他的N-K个点染成白色。将所有点染色后，你会获得黑点两两之间的距离加上白点两两之间距离的和的收益。

问收益最大值是多少。

输入

第一行两个整数N,K。

接下来N-1行每行三个正整数fr,to,dis，表示该树中存在一条长度为dis的边(fr,to)。

输入保证所有点之间是联通的。

N<=2000,0<=K<=N

输出

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

输入示例

5 2
1 2 3
1 5 1
2 3 1
2 4 2

输出示例

17

数据规模及约定

见“输入”

题解

树形 dp，设 f(i, j) 表示对于子树 i，选 j 个黑点所能得到的边的最大贡献，注意这里边的贡献是全局的，不是只考虑子树内部的，这样做比较方便转移。

那么在转移的时候，对于节点 u，我们在树上做背包，就是枚举每个儿子分别选几个黑点然后合并。这样看上去很暴力，但是我们可以证明总转移复杂度不超过 O(n²)。其实就是要证明 O(∑_u=1..n∑_{v,w是u的儿子}siz[v]·siz[w]) = O(n²)。我们不妨把“∑_{v,w是u的儿子}siz[v]·siz[w]”这一部分看做每次从 u 的每个子树中任选两个点 (a, b) 使得 a 和 b 的最近公公祖先恰好为 u 的方案数，那么每个节点只可能一次成为最近公公祖先，那么整个“∑_u=1..n∑_{v,w是u的儿子}siz[v]·siz[w]”其实就是树中点对个数，所以是 n² 的。

#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cctype>
#include <algorithm>
using namespace std;
#define LL long long

LL read() {
	LL x = 0, f = 1; char c = getchar();
	while(!isdigit(c)){ if(c == '-') f = -1; c = getchar(); }
	while(isdigit(c)){ x = x * 10 + c - '0'; c = getchar(); }
	return x * f;
}

#define maxn 2010
#define maxm 4010

int n, K, m, head[maxn], nxt[maxm], to[maxm];
LL dist[maxm];

void AddEdge(int a, int b, LL c) {
	to[++m] = b; dist[m] = c; nxt[m] = head[a]; head[a] = m;
	swap(a, b);
	to[++m] = b; dist[m] = c; nxt[m] = head[a]; head[a] = m;
	return ;
}

LL f[maxn][maxn];
int siz[maxn];
void build(int u, int fa) {
	siz[u] = 1;
	for(int e = head[u]; e; e = nxt[e]) if(to[e] != fa) {
		build(to[e], u);
		siz[u] += siz[to[e]];
	}
	if(siz[u] == 1){ f[u][0] = f[u][1] = 0; return ; }
	int sum = 1;
	f[u][0] = f[u][1] = 0;
	for(int e = head[u]; e; e = nxt[e]) if(to[e] != fa) {
		for(int k = sum; k >= 0; k--)
			for(int i = min(K, siz[to[e]]); i >= 0; i--) if(f[to[e]][i] >= 0 && k + i <= K)
				f[u][k+i] = max(f[u][k+i], f[to[e]][i] + f[u][k] + dist[e] * ((LL)i * (K - i) + ((LL)siz[to[e]] - i) * (n - K - siz[to[e]] + i)));
		sum = min(K, sum + siz[to[e]]);
	}
	return ;
}

int main() {
	n = read(); K = read();
	for(int i = 1; i < n; i++) {
		int a = read(), b = read(); LL c = read();
		AddEdge(a, b, c);
	}
	
	memset(f, -1, sizeof(f));
	build(1, 0);
	printf("%lld\n", f[1][K]);
	
	return 0;
}