树上染色

树上染色

树上背包 + 边贡献

HAOI2015] 树上染色 - 洛谷 | 计算机科学教育新生态 (luogu.com.cn)

#include <iostream>
#include <cstring>
#include <algorithm>
#include <vector>
#include <cmath>

using namespace std;
typedef long long ll;
const int N = 2e3 + 10;
const ll INF = 1e18;
int n, k;
ll f[N][N];//f[u][x]为以u为根的子树中，将x个点染成黑色
ll tmp[N];
int sz[N];
struct Edge
{
	int to;
	ll val;
};
vector<vector<Edge> > G(N);

void add(int u, int v, int w)
{
	G[u].push_back({v, w});
}

void dfs(int u, int fa)
{
	sz[u] = 1;
	for (auto [v, w] : G[u])
	{
		if (v == fa)
			continue;
		dfs(v, u);
		for (int i = 0; i <= sz[u] + sz[v]; i++)
			tmp[i] = -INF;
		for (int i = 0; i <= sz[u]; i++)
		{
			for (int j = 0; j <= sz[v]; j++)
			{
				ll cnt = j * (k - j) + (sz[v] - j) * (n - sz[v] - k + j);
				tmp[i+j] = max(tmp[i+j], f[u][i] + f[v][j] + w * cnt);
			}
		}
		for (int i = 0; i <= sz[u] + sz[v]; i++)
			f[u][i] = tmp[i];
		sz[u] += sz[v];
	}	
}
int main()
{
	ios::sync_with_stdio(false), cin.tie(0), cout.tie(0);
	cin >> n >> k;
	for (int i = 1; i < n; i++)
	{
		int u, v, w;
		cin >> u >> v >> w;
		add(u, v, w), add(v, u, w);
	}
	
	dfs(1, -1);
	
	cout << f[1][k] << endl;
	return 0;
}