题目描述
有一棵点数为 NNN 的树,树边有边权。给你一个在 0~N0~N0~N 之内的正整数 KKK,你要在这棵树中选择 KKK 个点,将其染成黑色,并将其他的 N−KN-KN−K 个点染成白色。 将所有点染色后,你会获得黑点两两之间的距离加上白点两两之间的距离的和的受益。问受益最大值是多少。0≤K≤N≤20000\le K\le N\le 20000≤K≤N≤2000。
算法分析
树形背包DP,设 f[i][j]f[i][j]f[i][j] 表示在以 iii 为根的子树中选择 jjj 个点染成黑色,子树内所有边对收益的最大贡献,根据子树内外黑白点的数目设计状态转移方程即可。
注意,下面写法中的 kkk 应倒序枚举,不然先枚举 000 的话,该子树内选择 000 个黑点所得的贡献也会被累加到该子树内选择更多黑点的答案中,从而使答案偏大。
代码实现
#include <cstdio>
#include <algorithm>
const int maxn=2005;
typedef long long int ll;
struct edge {int u,v,w;} edges[maxn<<1];
int head[maxn],nxt[maxn<<1],idx=0;
inline void add(int u,int v,int w) {
edges[++idx]=(edge){u,v,w};
nxt[idx]=head[u];head[u]=idx;
}
int n,K,tot[maxn];ll f[maxn][maxn];
void dfs(int x,int fa) {
tot[x]=1;
for(register int i=head[x];i;i=nxt[i]) {
edge &e=edges[i];if(e.v==fa) continue;
dfs(e.v,x);
for(register int j=std::min(tot[x],K);j>=0;--j) {
for(register int k=std::min(tot[e.v],K-j);k>=0;--k) {
int cnt=(K-k)*k+(n-K-tot[e.v]+k)*(tot[e.v]-k);
f[x][j+k]=std::max(f[x][j+k],f[x][j]+f[e.v][k]+(ll)cnt*e.w);
}
}
tot[x]+=tot[e.v];
}
}
int main() {
scanf("%d%d",&n,&K);
int m=n-1,u,v,w;idx=0;
while(m--) {
scanf("%d%d%d",&u,&v,&w);
add(u,v,w);add(v,u,w);
}
dfs(1,0);
printf("%lld\n",f[1][K]);
return 0;
}