【HAOI2015】树上染色

【HAOI2015】树上染色

这题思路好神仙啊，首先显然是树形dp，f[i][j]表示在以i为根的子树中选j个黑点对答案的贡献（并不是当前子树最大值），dp时只考虑i与儿子连边的贡献。此时（i，son[i]）产生的收益是（设子树大小为size[i]）子树上的黑点个数（j）与子树外的黑点个数（m - j）的乘积乘上这条边的边权（w[i]）加上子树上白点的个数（size[i] - j）乘以子树外白点的个数（n - m - size[i] + j）再乘以边权，这些贡献是在加入了根节点以后才产生的新的贡献，与子树上黑白点如何分配无关。

#include<iostream>
#include<cstring>
#include<cstdio>
#define int LL
#define LL long long
#define ma(x) memset(x,0,sizeof(x))
#define MAXN 3010
using namespace std;
struct edge
{
	int u,v,w,nxt;
	#define u(x) ed[x].u
	#define v(x) ed[x].v
	#define w(x) ed[x].w
	#define n(x) ed[x].nxt
}ed[2000000];
int first[MAXN],num_e;
#define f(x) first[x]
int f[MAXN][MAXN],du[MAXN],root;
int n,nk,size[MAXN],tmp[MAXN];
void dfs(int x,int fa)
{
	size[x]=1;
	for(int i=f(x);i;i=n(i))
	if(v(i)!=fa)dfs(v(i),x);
	for(int i=f(x);i;i=n(i))
	if(v(i)!=fa)
	{
		ma(tmp);
		for(int j=0;j<=size[x]&&j<=nk;j++)
			for(int k=0;k<=size[v(i)]&&k+j<=nk;k++)
			{
				tmp[j+k]=max(tmp[j+k],f[x][j]+f[v(i)][k]+k*(nk-k)*w(i)+(size[v(i)]-k)*(n-nk-size[v(i)]+k)*w(i));
			}
		for(int j=0;j<=nk;j++)
			f[x][j]=tmp[j];
		size[x]+=size[v(i)];
	}
}
inline void add(int u,int v,int w);
signed main()
{
//	freopen("in.txt","r",stdin);
//	freopen("2.in","r",stdin);

	scanf("%lld%lld",&n,&nk);
	int a,b,c;
	for(int i=1;i<n;i++)
	{
		scanf("%lld%lld%lld",&a,&b,&c);
		du[a]++;du[b]++;add(a,b,c);add(b,a,c);
	}
	for(int i=1;i<=n;i++)if(du[i]==1){root=i;break;}
	dfs(root,0);
	cout<<f[root][nk]<<endl;
}
inline void add(int u,int v,int w)
{
	++num_e;
	u(num_e)=u;
	v(num_e)=v;
	w(num_e)=w;
	n(num_e)=f(u);
	f(u)=num_e;
}