[BZOJ4033][HAOI2015]树上染色(树形DP)

4033: [HAOI2015]树上染色

Time Limit: 10 Sec  Memory Limit: 256 MB
Submit: 2437  Solved: 1034
[Submit][Status][Discuss]

Description

有一棵点数为N的树，树边有边权。给你一个在0~N之内的正整数K，你要在这棵树中选择K个点，将其染成黑色，并

将其他的N-K个点染成白色。将所有点染色后，你会获得黑点两两之间的距离加上白点两两之间距离的和的收益。

问收益最大值是多少。

Input

第一行两个整数N,K。

接下来N-1行每行三个正整数fr,to,dis，表示该树中存在一条长度为dis的边(fr,to)。

输入保证所有点之间是联通的。

N<=2000,0<=K<=N

Output

输出一个正整数，表示收益的最大值。

Sample Input

5 2
1 2 3
1 5 1
2 3 1
2 4 2

Sample Output

17
【样例解释】
将点1,2染黑就能获得最大收益。

HINT

2017.9.12新加数据一组 By GXZlegend

Source

鸣谢bhiaibogf提供

[Submit][Status][Discuss]

初看此题，树上背包！f[i][j]表示以i为根的子树选j个黑点的最大收益！

然后就发现转移爆炸。

于是令f[i][j]表示以i为根的子树选j个黑点的贡献，合并子树贡献后将自己到父亲的边的贡献加上即可，转移一句话。

然后就没了？

回头一看发现复杂度是$O(n^3)$的，于是分析一波复杂度。

$T(n)=\sum_{u=1}^{n}\sum_{v,w\ is\ a\ son\ of\ u} size[v]\times size[w]=\sum_{u=1}^{n}\sum_{u\ is\ the\ LCA\ of\ v,w} 1=O(n^2)$

于是就可做了，当然如果写丑了还是会变成$O(n^3)$的。

最后要记得边长是long long，狂WA不止。

#include<cstdio>
#include<cstring>
#include<algorithm>
#define rep(i,l,r) for (int i=l; i<=r; i++)
#define For(i,x) for (int i=h[x],k; i; i=nxt[i])
typedef long long ll;
using namespace std;

const int N=2010;
int n,m,u,v,w,cnt,h[N],sz[N],to[N<<1],nxt[N<<1];
ll f[N][N],dep[N],val[N<<1];
void add(int u,int v,int w){ to[++cnt]=v; val[cnt]=w; nxt[cnt]=h[u]; h[u]=cnt; }

void dfs(int x,int fa){
    f[x][0]=f[x][1]=0; sz[x]=1;
    For(i,x) if ((k=to[i])!=fa){
        dep[k]=val[i]; dfs(k,x);
        for (int l=min(sz[x],m); l>=0; l--)
            for (int j=min(sz[k],m-l); j>=0; j--) f[x][j+l]=max(f[x][j+l],f[x][l]+f[k][j]);
        sz[x]+=sz[k];
    }
    rep(i,0,min(sz[x],m)) f[x][i]+=dep[x]*(i*(m-i)+(sz[x]-i)*(n-sz[x]-m+i));
}

int main(){
    freopen("bzoj4033.in","r",stdin);
    freopen("bzoj4033.out","w",stdout);
    scanf("%d%d",&n,&m);
    memset(f,-0x3f,sizeof(f));
    rep(i,2,n) scanf("%d%d%d",&u,&v,&w),add(u,v,w),add(v,u,w);
    dfs(1,0); printf("%lld\n",f[1][m]);
    return 0;
}