题目链接
树形DP
简要题意
\(n\) 个点的树，其中 \(k\) 个点染黑色，\(n-k\) 个点染白色，求黑点两两距离之和加白点两两之和的最大值。
思路
我们首先考虑如果 \(k=0\)时，答案应该怎么算，此时显然是 \(\sum_{i=1}^n\sum_{j=i+1}^n dis(i,j)\)。
然后我们考虑如何在 \(O(n)\) 的时间复杂度内求出答案，明显我们可以将每条边的贡献拆开来算。
而经过边 \(u\to v\) 的次数明显为 \(size_v*(n-size_v)\)，很好理解，就是因为子树内每个点要到底子树外的点需要经过这条边。
这启发我们如何算黑白点的贡献，设 \(size1_x\) 表示以 \(x\) 为根的子树内黑点的数量，\(size2_x\) 表示以 \(x\) 为根的子树内白点的数量。
那么经过边 \(u\to v\) 的次数就是 \(size1_v*(k-size1_v)+size2_v*(n-k-size2_v)\)。此时就很好做了，设 \(dp_{i,j}\) 表示以 \(i\) 为根的子树内选 \(j\) 个黑点的边的最大贡献和。
那么对于边 \(u\to v\) 有
\[dp_{u,i+j}=\max_{i=1,j=1}^{k,i+j\le k} \{dp_{u,i}+dp_{v,j}+j*(k-j)*w_{u\to v}+(size_v-j)*(n-k-size_v+j)*w_{u\to v}\} \]\(Code\)
/*
*/
#include <bits/stdc++.h>
using namespace std;
inline long long read(){
long long x=0,w=0;char c=0;
while(!isdigit(c)) {w|=c=='-';c=getchar();}
while(isdigit(c)) x=(x<<3)+(x<<1)+(c^48),c=getchar();
return w?-x:x;
}
long long n,k;
struct edge{
int to,nex;
long long w;
}e[4050];
int h[2050],cnt;
inline void add(int x,int y,int w){
e[++cnt]={y,h[x],w};
h[x]=cnt;
}
long long dis[2050];
long long dp[2050][2050],siz[2050],f[2050];
void dfs(int x,int fa){
siz[x]=1;
for(int i=h[x];i;i=e[i].nex){
int v=e[i].to;
if(v==fa){
continue;
}
dis[v]=dis[x]+e[i].w;
dfs(v,x);
for(int i=0;i<=k;i++){
f[i]=0;
}
for(int l=0;l<=min(siz[x],k);l++){
for(int j=0;j<=min(siz[v],k);j++){
if(l+j>k){
break;
}
f[l+j]=max(f[l+j],dp[x][l]+dp[v][j]+e[i].w*j*(k-j)+e[i].w*(siz[v]-j)*(n-k-siz[v]+j));
}
}
siz[x]+=siz[v];
for(int j=0;j<=min(siz[x],k);j++){
dp[x][j]=f[j];
}
}
}
int main(){
n=read();
k=read();
for(int i=1,x,y,z;i<n;i++){
x=read();
y=read();
z=read();
add(x,y,z);
add(y,x,z);
}
dfs(1,0);
printf("%lld",dp[1][k]);
return 0;
}
/*
*/