【洛谷P3177】树上染色

题目

题目链接：https://www.luogu.com.cn/problem/P3177
有一棵点数为 \(n\) 的树，树边有边权。给你一个在 \(0 \sim n\) 之内的正整数 \(m\) ，你要在这棵树中选择 \(m\) 个点，将其染成黑色，并将其他 的 \(n-m\) 个点染成白色。将所有点染色后，你会获得黑点两两之间的距离加上白点两两之间的距离的和的受益。问受益最大值是多少。
\(n,m\leq 2000\)。

思路

显然树形 dp。设 \(f[x][i]\) 表示点 \(x\) 为根的子树内选择了 \(i\) 个黑点的最大贡献。
此时如果考虑分类 \(x\) 是黑色或白色再计算 \(x\) 的贡献显然是不可取的，因为点的贡献就涉及到其他相同颜色的点，而在 dp 状态中我们无法把其他点给加进来。
那么考虑计算每一条边的贡献。这样有一个好处：我们不需要知道这条边两边的点的具体位置，我们只需要知道两边分别有多少点，乘法原理计算即可。
考虑加入 \(x\) 的一棵子树 \(y\)，有转移

\[f[x][i+j]\gets \max\left (f'[x][i]+f[y][j]+\text{dis}(i,j)\times (j\times (m-j)+(\text{siz}[y]-j)\times (n-m-\text{siz}[y]+j))\right ) \]

发现这个转移第二维上界为子树大小，那么枚举到子树大小就可以做到 \(O(n^2)\)。

代码

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

const int N=2010;
int n,m,tot,head[N],siz[N];
ll f[N][N],g[N];

struct edge
{
	int next,to,dis;
}e[N*2];

void add(int from,int to,int dis)
{
	e[++tot]=(edge){head[from],to,dis};
	head[from]=tot;
}

void dfs(int x,int fa)
{
	siz[x]=1;
	for (int k=head[x];~k;k=e[k].next)
	{
		int v=e[k].to;
		if (v!=fa)
		{
			dfs(v,x);
			for (int i=0;i<=siz[x];i++)
				g[i]=f[x][i],f[x][i]=0;
			for (int i=0;i<=min(siz[x],m);i++)
				for (int j=0;j<=siz[v] && i+j<=m;j++)
					f[x][i+j]=max(f[x][i+j],g[i]+f[v][j]+e[k].dis*(1LL*j*(m-j)+1LL*(siz[v]-j)*(n-m-siz[v]+j)));
			siz[x]+=siz[v];
		}
	}
}

int main()
{
	memset(head,-1,sizeof(head));
	scanf("%d%d",&n,&m);
	for (int i=1,x,y,z;i<n;i++)
	{
		scanf("%d%d%d",&x,&y,&z);
		add(x,y,z); add(y,x,z);
	}
	dfs(1,0);
	printf("%lld",f[1][m]);
	return 0;
}