[树形DP]BZOJ 4033—— [HAOI2015]树上染色

题目梗概

给定一棵树，树边有边权。

将树上$K$个点染成黑色，其余$n-K$个点为白色。

使相同颜色节点的路径和最大。

解题思路

不难想到是一个树形DP的题目。

考虑每一条边要被计算的次数，我们在进行左儿子右兄弟的合并时，考虑一棵子树的根节点和其fa节点的合并，边就是连接这两个节点的边。

这条边加上的值=子树黑点数(n-子树黑点数)+子树白点数(k-子树白点数)

枚举子树黑点数，不断修正就可以了。

#include<cstdio>
#include<cstring>
#define LL long long
using namespace std;
const int maxn=2005,maxm=4005;
LL f[maxn][maxn],z[maxn];
int son[maxm],nxt[maxm],w[maxm],lnk[maxn],tot,s[maxn];
int n,k;
int add(int x,int y,int z){
    nxt[++tot]=lnk[x];lnk[x]=tot;son[tot]=y;w[tot]=z;
}
int min(int x,int y){if (x<y) return x;return y;}
LL max(LL x,LL y){if (x>y) return x;return y;}
void merge(int x,int y,int w){
    memset(z,0,sizeof(z));
    for (int i=0;i<=min(k,s[x]);i++)
    for (int j=0;j<=min(k-i,s[y]);j++)
    z[i+j]=max(z[i+j],f[x][i]+f[y][j]+(LL)w*((k-j)*j+(n-k-s[y]+j)*(s[y]-j)));
    for (int i=0;i<=k;i++) f[x][i]=max(f[x][i],z[i]);
}
void DFS(int x,int fa){
    s[x]=1;
    for (int j=lnk[x];j;j=nxt[j])if(son[j]!=fa){
        DFS(son[j],x);
        merge(x,son[j],w[j]);
        s[x]+=s[son[j]];
    }
}
int main(){
    scanf("%d%d",&n,&k);
    for (int i=1;i<n;i++){
        int x,y,z;
        scanf("%d%d%d",&x,&y,&z);
        add(x,y,z);add(y,x,z);
    }
    DFS(1,0);
    printf("%lld\n",f[1][k]);
    return 0;
}