树上染色(树形dp)

参考大佬题解:树上染色
原题：洛谷树上染色

/*********************************************************************
    程序名:
    版权: Joecai
    作者: Joecai
    日期: 2022-03-27 17:35
    说明:
*********************************************************************/
#include <bits/stdc++.h>
using namespace std;
#define x first
#define y second
# define rep(i,be,en) for(int i=be;i<=en;i++)
# define pre(i,be,en) for(int i=be;i>=en;i--)
typedef pair<int, int> PII;
#define ll long long
#define endl "\n"
#define LOCAL
#define pb push_back
const int N = 2500 + 10, INF = 0x3f3f3f3f;
ll qpow(int a, int b) {
	ll ans = 1;
	while (b) {
		if (b & 1) ans = ans * a;
		a = (ll)a * a;
		b >>= 1;
	}
	return ans;
}
inline ll read()
{
	ll X = 0;
	bool flag = 1;
	char ch = getchar();
	while (ch < '0' || ch > '9') {
		if (ch == '-') flag = 0;
		ch = getchar();
	}
	while (ch >= '0' && ch <= '9') {
		X = (X << 1) + (X << 3) + ch - '0';
		ch = getchar();
	}
	if (flag) return X;
	return ~(X - 1);
}

inline void write(ll X)
{
	if (X < 0) {
		X = ~(X - 1);
		putchar('-');
	}
	if (X > 9) write(X / 10);
	putchar(X % 10 + '0');
}
struct dsu {
	private:
		vector<int> par, ran, cnt;
	public:
		dsu(int n) {
			par.resize(n, 0);
			iota(par.begin(), par.end(), 0);
			ran.resize(n, 0);
			cnt.resize(n, 1);
		}
		int find(int x) {
			if (par[x] == x)return x;
			else return par[x] = find(par[x]);
		}
		bool unite(int x, int y) {
			x = find(x), y = find(y);
			if (x == y) return false;
			if (ran[x] < ran[y]) {
				par[x] = y;
				cnt[y] += cnt[x];
			}
			else {
				par[y] = x;
				cnt[x] += cnt[y];
				if (ran[x] == ran[y])ran[x]++;
			}
			return true;
		}
		bool same(int x, int y) {
			return find(x) == find(y);
		}
		int getcnt(int x) {
			return cnt[find(x)];
		}
};

struct node
{
	int to;
	int w;
};


vector<node>g[N];
int n, k;
ll f[N][N];// i 点i子树选j个点的贡献


int dfs(int u, int fa)
{
	int cnt = 1;
	f[u][0] = f[u][1] = 0;

	for (auto x : g[u])
	{
		int w = x.w;
		int j = x.to;
		if (j == fa) continue;
		int son = dfs(j, u);
		cnt += son;
		for (int i = min(k, cnt); i >= 0; i--)
			for (int r = 0; r <= min(i, son); r++)
			{
				//取r个
				if (f[u][i - r] != -1)
				{
					ll tot = (ll)r * (k - r) + (ll)(son - r) * (n - son - k + r);
					f[u][i] = max(f[u][i], f[j][r] + f[u][i - r] + 1ll * w * tot);
				}

			}

	}
	return cnt;
}



void solve()
{

	n = read();
	k = read();
	if(n-k<k)k=n-k;
	memset(f, -1, sizeof(f));
	for (int i = 1; i <= n - 1; i++)
	{
		int a, b, c;
		a = read();
		b = read();
		c = read();
		g[a].pb({b, c});
		g[b].pb({a, c});
	}

	dfs(1, -1);
	write(f[1][k]);
}




int main() {
	//#ifdef LOCAL
	//freopen("data.in.txt","r",stdin);
	//freopen("data.out.txt","w",stdout);
	//#endif
	int __ = 1;
	//cin>>__;
	while (__--)
	{
		solve();
	}
	return 0;
}

时间复杂度卡的死，不开O2优化100分，开O2accpeted.