洛谷P3177.树上染色

题目大意

一棵 $n(1\leq n\leq 2000)$ 个点的树，每条边有一个距离，从中选择 $k(0\leq k\leq n)$ 个点染成黑色，其余染成白色，最后我们可以得到黑色点两两之间的距离和加上白色点两两之间的距离和，求该值的最大值。

思路

我们考虑每一条边对答案的贡献，设该边为 $v$ 到 $to$ 的边，距离为 $d$ ，以 $to$ 为根的子树内有 $j$ 个黑点，于是其贡献为 $d(j(k-j)+(n-k-vsize[to]+j)(vsize[to]-j))$ ，于是我们设 $f[v,j]$ 为在以 $v$ 为根的子树中，有 $j$ 个点被染成黑色时，所有边对答案产生的贡献和的最大值，然后进行一个树形背包，最后 $f[0,k]$ 即为答案，复杂度 $O(n^2)$ 。

代码

#include<bits/stdc++.h>
#include<unordered_map>
#include<unordered_set>
using namespace std;
typedef long long LL;
typedef unsigned long long ULL;
typedef pair<int, int> PII;
#define all(x) x.begin(),x.end()
//#define int LL
//#define lc p*2+1
//#define rc p*2+2
#define endl '\n'
#define inf 0x3f3f3f3f
#define INF 0x3f3f3f3f3f3f3f3f
#pragma warning(disable : 4996)
#define IOS ios::sync_with_stdio(0),cin.tie(0),cout.tie(0)
const long double eps = 1e-10;
const long double pi = acos(-1.0);
const LL MOD = 100000000;
const LL mod = 998244353;
const int maxn = 2010;

LL N, M, f[maxn][maxn], vsize[maxn];

struct edge {
	LL to, cost;
};
vector<edge>G[maxn];

void add_edge(int from, int to, LL cost)
{
	G[from].push_back(edge{ to,cost });
	G[to].push_back(edge{ from,cost });
}

void dfs(int v, int p)
{
	vsize[v] = 1;
	f[v][0] = f[v][1] = 0;
	for (int i = 0; i < G[v].size(); i++)
	{
		edge& e = G[v][i];
		if (e.to == p)
			continue;
		dfs(e.to, v);
		for (LL j = min(vsize[v], M); j >= 0; j--)
		{
			for (LL k = min(vsize[e.to], M - j); k >= 0; k--)
				f[v][j + k] = max(f[v][j + k], f[v][j] + f[e.to][k] + ((M - k) * k + (N - M - vsize[e.to] + k) * (vsize[e.to] - k)) * e.cost);
		}
		vsize[v] += vsize[e.to];
	}
}

void solve()
{
	dfs(1, 0);
	cout << f[1][M] << endl;
}

int main()
{
	IOS;
	cin >> N >> M;
	int u, v, c;
	for (int i = 1; i < N; i++)
	{
		cin >> u >> v >> c;
		add_edge(u, v, c);
	}
	solve();

	return 0;
}