洛谷3177 [HAOI2015]树上染色

题意

将一棵 $N$ 个点的树上 $K$ 个点染黑，另外的点染白。求黑点两两之间距离和加上白点两两之间距离和的最大值。

$N,K\le 2000$ 。

题解

考虑每一条边的贡献，即为 边权 $\times ($ 一侧的黑点个数 $\times$ 另一侧的黑点个数 $+$ 一侧的白点个数 $\times$ 另一侧的白点个数 $)$ 。

用 $f[x][i]$ 表示以 $x$ 为根的子树，黑点有 $i$ 个的最大贡献。设上述边的贡献为 $cur$ ，方程式比较类似与树上背包：

$$f[x][i]=\max f[y][j]+f[x][i-j]+cur$$

注意转移之前应该先算出全为白点（即 $f[y][0]$）的情况，因为这样才构成与之后枚举等价的状态。

#include<bits/stdc++.h>
#define ll long long

using namespace std;

inline ll read()
{
    char ch=getchar(); ll f=1,x=0;
    while (ch<'0' || ch>'9') { if (ch=='-') f=-1; ch=getchar(); }
    while (ch>='0' && ch<='9') { x=x*10+ch-'0'; ch=getchar(); }
    return f*x;
}

struct Edge {
    ll next,to,w;
} edge[4005];
ll cnt,head[2005],n,m,a,b,c,f[2005][2005],siz[2005];

inline void add(ll u,ll v,ll w)
{
    edge[++cnt].to=v;
    edge[cnt].next=head[u];
    edge[cnt].w=w;
    head[u]=cnt;
}

void dfs(ll x,ll fa)
{
    siz[x]=1;
    f[x][0]=f[x][1]=0;
    for (ll i=head[x];i;i=edge[i].next)
    {
        ll y=edge[i].to,w=edge[i].w;
        if (y==fa) continue;
        dfs(y,x);
        siz[x]+=siz[y];
        for (ll j=min(siz[x],m);~j;j--)
        {
            if (~f[x][j]) f[x][j]+=f[y][0]+siz[y]*(n-m-siz[y])*w;
            for (ll k=min(siz[y],j);k;k--)
            {
                if (!~f[x][j-k]) continue;
                ll cur=k*(m-k)+(siz[y]-k)*(n-m-siz[y]+k); cur*=w;
                f[x][j]=max(f[x][j],f[y][k]+f[x][j-k]+cur);
            }
        }
    }
}

signed main()
{
    n=read(); m=read();
    for (ll i=1;i<n;i++)
    {
        a=read(); b=read(); c=read();
        add(a,b,c);
        add(b,a,c);
    }
    memset(f,-1,sizeof(f));
    dfs(1,0);
    return !printf("%lld",f[1][m]);
}