P3177 [HAOI2015] 树上染色

题目传送门

树形背包

题意

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

题解

首先套路的设 \(f_{u,t}\) 为 \(u\) 子树内选 \(t\) 个黑色点的答案，并套路的做树上背包。

一开始想的比较朴素，直接以点对为对象统计答案，但是这样搞得话合并子树的时候还要记录子树内到根节点的距离和，要多维护一个数组，感觉很难写，所以就看看有没有更简单的方法。

刚才的问题是在于我们必须关注每个点在子树中的位置，既然统计的是边的贡献，何妨不以边为对象，去关注一条边对答案的贡献呢？

对于一条边，我们只需要知道它被经过了几次，就能算出它产生的贡献，显然，边的两侧有多少组同色的点，边就会被经过几次，于是乎，我们只需要枚举子树内染成黑色的点的数量，就能算出合并时连接两颗树的边的贡献了。

当 \(v\) 子树要合并到 \(u\) 子树中，并且 \(v\) 子树中有 \(j\) 个黑色点时，连接两颗树的边被经过的次数为 \(j(m-j)+(sz_v-j)(n-sz_v-(m-j))\)。

转移显然。

代码

#include <bits/stdc++.h>
#define int long long
using namespace std;
const int N=2005;
int n,k;
int f[N][N],sz[N],g[N];
vector<pair<int,int>> e[N];void dfs(int u,int fa){sz[u]=1;f[u][0]=0;memset(g,0,sizeof g);for(const auto &it:e[u]){int v=it.first,w=it.second;if(v==fa) continue;dfs(v,u);for(int i=0;i<=min(sz[u],k);i++)for(int j=0;j<=min(sz[v],k-i);j++)g[i+j]=max(g[i+j],f[u][i]+f[v][j]+w*(j*(k-j)+(sz[v]-j)*(n-sz[v]-k+j)));sz[u]+=sz[v];for(int i=0;i<=sz[u];i++) f[u][i]=g[i];}
}signed main(){cin.tie(nullptr)->sync_with_stdio(false);cin>>n>>k;for(int i=1;i<n;i++){int u,v,w;cin>>u>>v>>w;e[u].emplace_back(v,w);e[v].emplace_back(u,w);}dfs(1,1);cout<<f[1][k];return 0;
}