[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$ 的次数明显为 $siz{e}_v\ast (n-siz{e}_v)$，很好理解，就是因为子树内每个点要到底子树外的点需要经过这条边。

这启发我们如何算黑白点的贡献，设 $size{1}_x$ 表示以 $x$ 为根的子树内黑点的数量，$size{2}_x$ 表示以 $x$ 为根的子树内白点的数量。
那么经过边 $u\to v$ 的次数就是 $size{1}_v\ast (k-size{1}_v)+size{2}_v\ast (n-k-size{2}_v)$。此时就很好做了，设 $d{p}_{i,j}$ 表示以 $i$ 为根的子树内选 $j$ 个黑点的边的最大贡献和。

那么对于边 $u\to v$ 有

$$d{p}_{u,i+j}=\underset{i=1,j=1}{\overset{k,i+j\le k}{max}}\{d{p}_{u,i}+d{p}_{v,j}+j\ast (k-j)\ast w_{u\to v}+(siz{e}_v-j)\ast (n-k-siz{e}_v+j)\ast 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; 
}
/*

*/