[HAOI2015] 树上染色

题目链接

树形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; 
}
/*

*/