【HAOI 2015】树上染色

题目描述

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

算法分析

树形背包DP，设 $f[i][j]$ 表示在以 $i$ 为根的子树中选择 $j$ 个点染成黑色，子树内所有边对收益的最大贡献，根据子树内外黑白点的数目设计状态转移方程即可。

注意，下面写法中的 $k$ 应倒序枚举，不然先枚举 $0$ 的话，该子树内选择 $0$ 个黑点所得的贡献也会被累加到该子树内选择更多黑点的答案中，从而使答案偏大。

代码实现

#include <cstdio>
#include <algorithm>
const int maxn=2005;
typedef long long int ll;
struct edge {int u,v,w;} edges[maxn<<1];
int head[maxn],nxt[maxn<<1],idx=0;
inline void add(int u,int v,int w) {
	edges[++idx]=(edge){u,v,w};
	nxt[idx]=head[u];head[u]=idx;
}
int n,K,tot[maxn];ll f[maxn][maxn];
void dfs(int x,int fa) {
	tot[x]=1;
	for(register int i=head[x];i;i=nxt[i]) {
		edge &e=edges[i];if(e.v==fa) continue;
		dfs(e.v,x);
		for(register int j=std::min(tot[x],K);j>=0;--j) {
			for(register int k=std::min(tot[e.v],K-j);k>=0;--k) {
				int cnt=(K-k)*k+(n-K-tot[e.v]+k)*(tot[e.v]-k);
				f[x][j+k]=std::max(f[x][j+k],f[x][j]+f[e.v][k]+(ll)cnt*e.w);
			}
		}
		tot[x]+=tot[e.v];
	}
}
int main() {
	scanf("%d%d",&n,&K);
	int m=n-1,u,v,w;idx=0;
	while(m--) {
		scanf("%d%d%d",&u,&v,&w);
		add(u,v,w);add(v,u,w);
	}
	dfs(1,0);
	printf("%lld\n",f[1][K]);
	return 0;
}