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

[HAOI2015]树上染色

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