[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

#include<bits/stdc++.h>
#define ll long long
using namespace std;

const int maxn = 2e3+10;

int n,m,size,head[maxn],sum[maxn];

ll f[maxn][maxn];

struct edge {
    int v,w,nex;
}e[maxn<<1];

void adde(int u,int v,int w) {
    e[size].v=v;e[size].w=w;e[size].nex=head[u];head[u]=size++;
}

void dfs(int u,int fa) {
    f[u][1]=f[u][0]=0;sum[u]=1;
    for(int i=head[u];~i;i=e[i].nex) {
        int v=e[i].v;
        if(v==fa) continue;
        dfs(v,u);
        sum[u]+=sum[v];
        for(int j=min(m,sum[u]);j>=0;j--) {
            for(int k=0;k<=min(sum[v],j);k++) {
                if(f[u][j-k]==-1) continue;
                ll val = (ll) ((ll)k*(m-k)+(ll)(sum[v]-k)*(n-m-sum[v]+k))*e[i].w;
                f[u][j]=max(f[u][j],f[u][j-k]+f[v][k]+val);
            }
        }
    }
}

int main() {
    freopen("in.txt","r",stdin);
    memset(head,-1,sizeof(head));
    scanf("%d%d",&n,&m);
    if(m+m>n) m=n-m;
    for(int i=1;i<=n-1;i++)  {
        int u,v,w;scanf("%d%d%d",&u,&v,&w);
        adde(u,v,w);adde(v,u,w);
    }
    memset(f,-1,sizeof(f));
    dfs(1,-1);
    printf("%lld",f[1][m]);
    return 0;
}