[HAOI2015]树上染色（树形dp）

[HAOI2015]树上染色

题目描述

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

输入输出格式

输入格式：

第一行包含两个整数 N, K 。接下来 N-1 行每行三个正整数 fr, to, dis ， 表示该树中存在一条长度为 dis 的边 (fr, to) 。输入保证所有点之间是联通的。

输出格式：

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

输入输出样例

输入样例#1： 复制

3 1
1 2 1
1 3 2

输出样例#1： 复制

3

说明

对于 100% 的数据， 0<=K<=N <=2000

题解

最开始我以为要处理出点与点之间的距离。
然后对于k的话实际上就是min(k,n-k)。
然后dp出最小价值的k个点对。
拿总路径和减去最小dp值。
但是发现不好维护。

于是抄看了题解
对于一个子树内我要选取的黑点。
我们这一次dp的不仅是增加的黑点的价值，还要处理出减少的白点的价值。
也就是说每选一个点，就要判断这条路径的贡献变化了多少。
对于每条路径的贡献。
为当前子树的黑节点×子树外的黑节点×边权+当前子树的白节点×子树外的白节点×边权就可以了。
这样就不用刻意去记录点对了。

代码

#include<cstdio>
#include<cstring>
#include<cmath>
#include<iostream>
#include<algorithm>
#define ll long long
using namespace std;
const int N=2001;
struct node{
    int to,nex;
    ll v;
}e[N<<2];
int n,k,num,head[N],size[N];
ll f[N][N];
ll read(){
    ll x=0,w=1;char ch=getchar();
    while(ch>'9'||ch<'0'){if(ch=='-')w=-1;ch=getchar();}
    while(ch>='0'&&ch<='9')x=x*10+ch-'0',ch=getchar();
    return x*w;
}

void add(int from,int to,int v){
    num++;
    e[num].to=to;
    e[num].v=v;
    e[num].nex=head[from];
    head[from]=num;
}

void dfs1(int x,int fa){
    size[x]=1;
    for(int i=head[x];i;i=e[i].nex){
        int v=e[i].to;
        if(v==fa)continue;
        dfs1(v,x);size[x]+=size[v];
    }
}

void dfs2(int x,int fa){
    f[x][0]=0;f[x][1]=0;
    for(int i=head[x];i;i=e[i].nex){
        int v=e[i].to;if(v!=fa){dfs2(v,x);
        for(int j=min(size[x],k);j>=0;j--){
            for(int l=0;l<=min(size[v],j);l++)
            if(f[x][j-l]!=-1){
            ll val=1ll*(l)*(k-l)*e[i].v+1ll*(size[v]-l)*(n-k+l-size[v])*e[i].v;
            f[x][j]=max(f[x][j-l]+f[v][l]+val,f[x][j]);
            }
            }
        }
    }
}

int main(){
    n=read();k=read();
    for(int i=1;i<n;i++){
        int  x=read(),y=read(),z=read();
        add(x,y,z);add(y,x,z);
    }
    memset(f,-1,sizeof(f));
    dfs1(1,1);dfs2(1,1);
    printf("%lld\n",f[1][k]);
    return 0;
}