[bzoj4033]树上染色

题目描述

有一棵点数为N的树，树边有边权。给你一个在0~N之内的正整数K，你要在这棵树中选择K个点，将其染成黑色，并
将其他的N-K个点染成白色。将所有点染色后，你会获得黑点两两之间的距离加上白点两两之间距离的和的收益。
问收益最大值是多少。

DP

首先贡献很恶心，但我们可以这样考虑，只考虑每条边的贡献。
设f[i,j]表示以i为根的子树中染黑了j个点的最大贡献，这里只考虑了子树中每条边以及i的父亲边的贡献。
那么转移很显然。
复杂度呢？n^3？
我们合并并不需要每次枚举m那么多，实际上枚举量等于子树大小。
那么我们改改枚举，复杂度就是n^2了！
因为可以注意到任意两点会在其lca处被计算，实际上枚举量只有n^2。

#include<cstdio>
#include<algorithm>
#define fo(i,a,b) for(i=a;i<=b;i++)
#define fd(i,a,b) for(i=a;i>=b;i--)
#define min(a,b) (a<b?a:b)
using namespace std;
typedef long long ll;
const int maxn=2000+10;
ll f[maxn][maxn],g[maxn];
int h[maxn],go[maxn*2],dis[maxn*2],nxt[maxn*2],size[maxn];
int i,j,k,l,t,n,m,tot;
void add(int x,int y,int z){
    go[++tot]=y;
    dis[tot]=z;
    nxt[tot]=h[x];
    h[x]=tot;
}
void dfs(int x,int y,int z){
    int i,j,k,t=h[x];
    while (t){
        if (go[t]!=y){
            dfs(go[t],x,dis[t]);
            fo(i,0,m) g[i]=0;
            fo(i,0,min(size[x],m))
                fo(j,0,min(size[go[t]],m-i))
                    g[i+j]=max(g[i+j],f[x][i]+f[go[t]][j]);
            fo(i,0,m) f[x][i]=g[i];
            size[x]+=size[go[t]];
        }
        t=nxt[t];
    }
    fd(i,m-1,0) f[x][i+1]=max(f[x][i+1],f[x][i]);
    size[x]++;
    fo(i,0,m){
        f[x][i]+=(ll)z*i*(m-i);
        f[x][i]+=(ll)z*(size[x]-i)*(n-m-size[x]+i);
    }
}
int main(){
    scanf("%d%d",&n,&m);
    fo(i,1,n-1){
        scanf("%d%d%d",&j,&k,&l);
        add(j,k,l);add(k,j,l);
    }
    dfs(1,0,0);
    printf("%lld\n",f[1][m]);
}