[洛谷-P3177][HAOI2015] 树上染色（树形DP + 背包DP）

一、问题

[HAOI2015] 树上染色

题目描述

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

输入格式

第一行包含两个整数 $n,k$。

第二到 $n$ 行每行三个正整数 $u,v,w$，表示该树中存在一条长度为 $w$ 的边 $(u,b)$。输入保证所有点之间是联通的。

输出格式

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

样例 #1

样例输入 #1

3 1
1 2 1
1 3 2

样例输出 #1

3

提示

对于 $100\%$ 的数据，$0\le n,k\le 2000$。

二、思路

1、贡献分析

简单的来说就是给我们一个无根的无向带权的树，这棵树上的节点一开始都是白色的，然后我们要将其中的$k$个点染成黑色，然后算出相同颜色的节点之间的距离，然后把所有的距离加在一起，我们的目的就是求出所有距离和的最大值。

那么这道题怎么做呢？我们看下面的图：

我们以橙色边为研究对象，橙色边的边权记作：$w$。观察这条边对于我们最终答案的贡献。

我们将$son$所在的子树（包括$son$）的节点总数记作$siz(son)$。那么如果我们在子树中选择$q$个节点染成黑色的话，就会剩下$siz(son)-q$个白色的节点。由于我们在整棵树中一共要选择$k$个节点染成黑色，所以在蓝色区域（即除了$son$所在子树的其余点）中就会有$k-q$个黑色节点。那么当我们计算这条边在黑色点之间的贡献时，这条边会被利用$q\ast (k-q)$次。那么对于黑色点之间的路径和的贡献就是$w\ast q\ast (k-q)$。

那么白色点我们如何计算呢？由于我们一共有$n$个点，那么除去$k$个黑色点，我们还剩下$n-k$个白色节点，由于我们的红色区域（即$son$所在的子树）中有$siz(son)-q$个白色节点，则我们的蓝色区域内则有$n-k-(siz(son)-q)$个白色节点。那么这条边对于白色点的路径和的贡献就是：$w\ast [siz(son)-q]\ast [n-k-(siz(son)-q)]$

综合上面两段的叙述，我们这条边对于答案的总贡献就是：$$w\ast q\ast (k-q)+w\ast [siz(son)-q]\ast [n-k-(siz(son)-q)]$$

剩下的过程其实就是树上背包了。如果大家不懂的话，建议先去看一看作者之前的文章：

AcWing 10. 有依赖的背包问题（分组背包问题 + 树形DP）

2、状态表示

$f[u][j]$表示在以$u$为根的树中（包括$u$）选择$k$个点染成黑色点的条件下， 相同颜色点之间的路径和的最大值。

3、状态转移

根据我们刚才的图：我们发现图中的红色区域可以用$f[son][q]$表示，蓝色区域可以用$f[u][j-q]$表示。同时我们还需要加上这条橙色边贡献。

$$f[u][j]=max(f[u][j],f[son][q]+f[u][j-q]+w\ast q\ast (k-q)+w\ast [siz(son)-q]\ast [n-k-(siz(son)-q)])$$

4、初末状态

$f[u][0]$表示在以$u$为根的树中，选0个点染成黑色，那么最大值明显是$0$。
$f[u][1]$表示在以$u$为根的树中，选1个点染成黑色，由于两点之间才存在路径，所以$1$个点的时候，最大值也是$0$。

这道题还存在很多不合法的状态，什么意思呢？
我们的转移方程中用到了$f[u][j-q]$，如果我们以$u$为根的树的节点个数（包括$u$）小于$j-q$的话，这个状态就是不合法的。如果我们不把这些状态标记出来的话，我们的状态表示就会因此发生变化，什么变化呢？

状态表示会从恰好有$k$个黑色点转变成不超过$k$个黑色点。

如若状态表示发生变化的话，我们的答案就无法保证正确性了。

同时，我们通过特判这些不合法的状态也可以减少我们的时间损耗。所以我们需要将不合法的状态初始化为$-1$。

当我们在写转移方程的时候，如果用到了某个状态的值是$-1$的话，就说明这个状态是不合法的。为什么呢？
因为我们的DP一定是去不重不漏地枚举了所有子问题状态，如果某个状态是$-1$的话，就说明这是个违法的状态，所以才没有被更新。

三、代码

#include<bits/stdc++.h>
#define endl '\n'
#define INF 0x3f3f3f3f
using namespace std;
typedef long long ll;
typedef pair<int,int> pii;
const int N = 2e3 + 10;
vector<pii>edge[N];
ll n, k;
ll siz[N], f[N][N];

int cal_son(int u, int father)
{
	int sum = 0;
	for(int i = 0; i < edge[u].size(); i ++ )
	{
		int son = edge[u][i].first;
		if(son == father)
			continue;
		sum += cal_son(son, u);
	}
	siz[u] = sum + 1;
	return sum + 1;
}


void dp(int u, int father)
{
	f[u][0] = f[u][1] = 0;
	for(int i = 0; i < edge[u].size(); i ++)
	{
		int son = edge[u][i].first;
		ll w = edge[u][i].second;
		
		if(son == father)
			continue;

		dp(son, u);

		for(int j = min(k, siz[u]); j >= 0; j -- )
		{
			for(int q = 0; q <= min(siz[son], (ll)j); q ++ )
			{
				if(f[u][j - q] == -1)
					continue;
				ll tot = (ll)(q) * (k - q) + (ll)(siz[son] - q) * (n - k - siz[son] + q);
				// cout << tot << endl;

				ll sum = (ll)tot * w;
				// cout << sum << endl;
				f[u][j] = max(f[u][j - q] + f[son][q] + sum, f[u][j]);
			}
		}
	}
}
/*
	f[u][i]：以u为根的子树中，选择i个点涂成黑色，答案的最大值,不包括u点
	f[u][i] = max: f[u][i - k] + f[son][k] + sum
*/
void solve()
{
	cin >> n >> k;
	memset(f, -1, sizeof f);
	for(int i = 0; i < n - 1; i ++ )
	{
		int a, b , c;
		cin >> a >> b >> c;
		edge[a].push_back({b, c});
		edge[b].push_back({a, c});
	}

	cal_son(1, 0);
	dp(1, 0);

	cout << f[1][k] << endl;
}

int main()
{
	ios::sync_with_stdio(0);
	cin.tie(0);
	cout.tie(0);

	solve();
}