【刷题笔记】树形 DP

树上染色

题解

容易想到设 \(f_{u,i}\) 表示子树内染了 \(i\) 个黑点的最大答案，但是发现对子树外节点的贡献不好统计，所以将这条边对子树外节点的贡献 提前 计算进去。

DP 转移：
\(f_{u,i}=\max(f_{u,i-j}+f_{v,j}+w_{u,v}\times j\times(k-j)+w_{u,v}\times (sz_v-j)\times (n-sz_v-k+j))\)
DP 解释：子树 \(v\) 内有 \(j\) 个黑点，\(sz_v - j\) 个白点，子树外有 \(k-j\) 个黑点，\(n-sz_v-k+j\) 个白点，两两相乘计算即可。

但是注意到这样总时间复杂度是 \(O(nk^2)\) 的，所以考虑优化（背包优化）：限定枚举 \(i,j\) 的上界，其中 \(i \le \min(sz_u, k),\max(i - sz_u + sz_v,0) \le j\le \min(i,sz_v)\)（一般只要在树上背包时不多算，时间复杂度就是 \(O(n^2)\) 的）

code

#include<bits/stdc++.h>
#define int long long
#define N 3010
#define pii pair<int, int>
#define fi first 
#define se second
#define pb push_back
#define mkp make_pair
#define Fo(a, b) for(auto a : b)
#define fo(a, b, c) for(int b = a; b <= c; b++)
#define _fo(a, b, c) for(int b = a; b >= c; b--)
using namespace std;
int n, k, sz[N], f[N][N], g[N];
vector<pii>G[N];
void dp(int u, int fa){
	sz[u] = 1;
	Fo(v, G[u]){
		if(v.fi == fa) continue;
		dp(v.fi, u), sz[u] += sz[v.fi];// 必须边做边求 
		_fo(min(sz[u], k), i, 0) fo(max(0ll, i - sz[u] + sz[v.fi]), j, min(i, sz[v.fi])){
			f[u][i] = max(f[u][i], f[u][i - j] + f[v.fi][j] + v.se * j * (k - j) + v.se * (sz[v.fi] - j) * (n - sz[v.fi] - k + j));
		}
	}
}
signed main(){
	ios::sync_with_stdio(0);
	cin.tie(0); cout.tie(0);
	cin >> n >> k;
	fo(1, i, n - 1){
		int x, y, w; cin >> x >> y >> w;
		G[x].pb(mkp(y, w)), G[y].pb(mkp(x, w));
	}
	dp(1, 0);
	cout << f[1][k];
	return 0;
}