luogu P3177 [HAOI2015]树上染色 |树形DP

题目描述

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

输入格式

第一行包含两个整数 (n,k)。

第二到 (n) 行每行三个正整数 (fr, to, dis)，表示该树中存在一条长度为 (dis) 的边 ((fr, to))。输入保证所有点之间是联通的。

输出格式

输出一个正整数，表示收益的最大值。

---

转移方程为dp[u][i] = max( dp[u][i], dp[u][i-j] + dp[v][j] + val )

其中v为u的子节点，j为在这个子节点中选择的黑色点的个数，val为这条边的贡献

val = j(k-j)w + (sz[v]-j)(n-k+j-sz[v])w

其中w为这条边的边权，n为总的节点数，k为总的需要选择的黑色节点数，sz[v]为以v为根的子树的节点数量

为了卡时限，可以尝试rand一个根,~不过也有可能会变慢~

---

#include<cstdio>
#include<cstring>
#include<cstdlib>
#include<set>
#include<ctime>
#include<queue>
#include<iostream>
#include<cmath>
#include<algorithm>
#define inf 1000000000
#define ll long long
#define int long long
using namespace std;
inline int read(){
	int x=0,f=1;char ch=getchar();
	while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
	while(ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
	return x*f;
}
const int N=2005;
int nxt[N<<1],head[N],go[N<<1],cost[N<<1],tot;
inline void add(int u,int v,int w){
	nxt[++tot]=head[u],head[u]=tot,go[tot]=v,cost[tot]=w;
	nxt[++tot]=head[v],head[v]=tot,go[tot]=u,cost[tot]=w;
}
int n,m,dp[N][N],siz[N];
inline void DP(int &u,int &v,int &w){
	for(int i=min(m,siz[u]);i>=0;i--)
	for(int j=0;j<=min(i,siz[v]);j++)
	if(~dp[u][i-j]){
		int val=w*(j*(m-j)+(siz[v]-j)*(n-m+j-siz[v]));
		dp[u][i]=max(dp[u][i],dp[u][i-j]+dp[v][j]+val);
	}
}
void dfs(int u,int fa){
	siz[u]=1,dp[u][0]=dp[u][1]=0;
	for(int i=head[u];i;i=nxt[i]){
		int v=go[i],w=cost[i];
		if(v==fa)continue;
		dfs(v,u);
		siz[u]+=siz[v];
		DP(u,v,w);
	}
}
signed main(){
	//srand(time(0));
	memset(dp,-1,sizeof(dp));
	n=read(),m=read();
	m=min(m,n-m);
	for(int i=1,u,v,w;i<n;i++){
		u=read(),v=read(),w=read();
		add(u,v,w);
	}
	int rt=1;
	dfs(rt,0);
	cout<<dp[rt][m]<<endl;
}