[HAOI2015]树上染色(树上dp)

[HAOI2015]树上染色

这种要算点对之间路径的长度和的题,难以统计每个点的贡献.这个时候一般考虑算每一条边贡献了哪些点对.
知道这个套路以后,那么这题就很好做了.
状态:设\(dp[u][i]\)表示u节点(子树里有i个黑点)的子树的边的贡献的和.
转移:转移就很好想了,知道v内的黑点个数j,知道v内的白点数目\(sz[v]-j\),知道总共的黑点数目\(m\),知道总共的白点数目\((n-m)\),知道边权w,那么转移方程显然就是:

答案就是\(dp[1][k]\)

#include<bits/stdc++.h>
#define maxn 2005
#define ll long long
using namespace std;
int n,m,sz[maxn];
ll dp[maxn][maxn],t[maxn];
int head[maxn],nxt[maxn<<1],to[maxn<<1],w[maxn<<1],cnt;
void add(int u,int v,int ww)
{
	nxt[++cnt]=head[u];head[u]=cnt;
	w[cnt]=ww;to[cnt]=v;
}
void dfs(int u,int ff)
{
	dp[u][0]=dp[u][1]=0;sz[u]=1;
	for(int ii=head[u];ii;ii=nxt[ii])
	{
		int v=to[ii];if(v==ff)continue;
		dfs(v,u);sz[u]+=sz[v];
		for(int i=0;i<=m;i++)t[i]=dp[u][i];
		for(int i=m;i>=0;i--)
		{
			for(int j=0;j<=min(i,sz[v]);j++)
			{
				if(t[i-j]==-1)continue;ll res=0,ww=w[ii];
				res=dp[v][j]+ww*(m-j)*j+ww*(sz[v]-j)*(n-m-sz[v]+j);
				//cout<<u<<" "<<v<<" "<<ww<<" "<<i<<" "<<j<<endl;
				dp[u][i]=max(dp[u][i],res+t[i-j]);
			}
		}
		//cout<<u<<"&"<<v<<" ";
		//for(int i=0;i<=m;i++)cout<<dp[u][i]<<" ";cout<<endl;
	}
}
int main()
{
	cin>>n>>m;
	for(int i=1,u,v,ww;i<n;i++)
		scanf("%d%d%d",&u,&v,&ww),add(u,v,ww),add(v,u,ww);
	memset(dp,-1,sizeof(dp));
	dfs(1,0);cout<<dp[1][m]<<endl;
	//for(int u=1;u<=n;u++,cout<<endl)
	//	for(int i=0;i<=m;i++)
	//		cout<<dp[u][i]<<" ";
	return 0;
}