[Luogu] P3177 [HAOI2015]树上染色

\(Link\)

Description

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

Solution

妙妙的树形\(DP\)。

我们设\(dp[x][t]\)表示在以\(x\)为根的子树中选择\(t\)个点染黑时的最大贡献。

发现距离不太好搞，那么我们可以将距离转化为路径，然后再将路径拆分成边，就可以记录每条边被经过的次数，直接计算即可。

我们枚举\(x\)的子节点\(y\)，先处理\(y\)这棵子树和其他子树，再对\(x\)做贡献。分别枚举\(y\)中被染黑的节点数\(tn\)和\(x\)中除了\(y\)的子树被染黑的节点数\(te\)。注意要倒序枚举，避免后效性。那么这样产生的贡献就是\(value=z*(k - tn) * tn + (sz[y] - tn) * (n - k - sz[y] + tn)\)（黑点贡献加上白点贡献）。

总的方程就是\(dp[x][te+tn]=max\{dp[x][te]+dp[y][tn]+value\}\)。

注意枚举\(sz[]\)的时候，不要在最开始就一遍算出来，要一边算，一遍处理\(DP\)数组。这样就不会重复算，防止超时。因为\(x\)可能会有多棵子树，我们一棵一棵枚举，现在的\(dp[x][te]\)其实就是之前的\(dp[x][te+tn]\)，正好已经算过。

Code

#include <bits/stdc++.h>

using namespace std;

#define ll long long

int n, k, tot, hd[2005], to[4005], nxt[4005], w[4005], sz[2005];

ll dp[2005][2005];

int read()
{
	int x = 0, fl = 1; char ch = getchar();
	while (ch < '0' || ch > '9') { if (ch == '-') fl = -1; ch = getchar();}
	while (ch >= '0' && ch <= '9') {x = x * 10 + ch - '0'; ch = getchar();}
	return x * fl;
}

void add(int x, int y, int z)
{
	tot ++ ;
	to[tot] = y;
	w[tot] = z;
	nxt[tot] = hd[x];
	hd[x] = tot;
	return;
}

void dfs(int x, int fa)
{
	sz[x] = 1;
	for (int i = hd[x]; i; i = nxt[i])
	{
		int y = to[i], z = w[i];
		if (y == fa) continue;
		dfs(y, x);
		sz[x] += sz[y];
		for (int te = sz[x] - sz[y]; te >= 0; -- te) //　注意是倒序枚举！
			for (int tn = sz[y]; tn >= 0; -- tn)
				dp[x][te + tn] = max(dp[x][te + tn], dp[x][te] + dp[y][tn] + 1ll * z * ((k - tn) * tn + (sz[y] - tn) * (n - k - sz[y] + tn)));
	}
	return;
}

int main()
{
	n = read(); k = read();
	for (int i = 1; i <= n - 1; i ++ )
	{
		int x = read(), y = read(), z = read();
		add(x, y, z); add(y, x, z);
	}
	dfs(1, 0);
	printf("%lld\n", dp[1][k]);
	return 0;
}