洛谷 3177 [HAOI2015] 树上染色

题目描述

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

输入输出格式

输入格式：

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

输出格式：

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

输入输出样例

输入样例#1：

输出样例#1：

说明

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

题解：

最终的收益等于每一条边的收益（权值乘以被用到的次数）的和，假设dp[i][j]表示以i为根的子树内有j个点为黑点时边的收益，从叶子到根的顺序计算每条边的收益，遍历到i点时，i的子树里的边已经被计算过了，i到其儿子的边被新加了进来，用这些新加的边的收益和子树的收益更新以i为根的子树的收益。

转移方程为dp[i][j]=dp[k][l]+边权*子树内黑点*子树外黑点+边权*子树内白点*子树外白点。

一开始我以为自己计算的是答案的两倍，这样智障了很久，后来把/2删掉后过了，我又仔细思考了一下，我是一条边一条边的算的，所以只会算一次。

#include<algorithm>

#include<iostream>

#include<cstdlib>

#include<cstdio>

#define nn 2010

#define mm 4010

#define lo long long

#define inf -100000000

using namespace std;

int e=0;

int fir[nn],nxt[mm],to[mm],w[mm],size[nn],n,k;

lo dp[nn][nn];

bool vis[nn];

int get()

{

    int ans=0,f=1;char ch=getchar();

    while(!isdigit(ch)) {if(ch=='-') f=-1;ch=getchar();}

    while(isdigit(ch)) {ans=ans*10+ch-'0';ch=getchar();}

    return ans*f;

}

void add(int a,int b,int c)

{

    nxt[++e]=fir[a];fir[a]=e;to[e]=b;w[e]=c;

    nxt[++e]=fir[b];fir[b]=e;to[e]=a;w[e]=c;

}

void dfs(int o)

{

    size[o]=1;

    for(int i=fir[o];i;i=nxt[i])

      if(!vis[to[i]])

      {

          vis[to[i]]=1;

          dfs(to[i]);

          size[o]+=size[to[i]];

      }

}

void solve(int o)

{

    dp[o][0]=0;

    dp[o][1]=0;

    for(int i=fir[o];i;i=nxt[i])

      if(size[o]>size[to[i]])

      {

          solve(to[i]);

          for(int j=size[o];j>=0;j--)                //把size[o]写成了size[i]

            for(int p=0;p<=size[to[i]]&&p<=j;p++)

              dp[o][j]=max(dp[o][j],dp[to[i]][p]+dp[o][j-p]+(lo)p*(k-p)*w[i]+(lo)w[i]*(size[to[i]]-p)*(n-k-size[to[i]]+p));

      }

}

int main()

{

    n=get(),k=get();

    int a,b,c;

    for(int i=1;i<n;i++)

    {

        a=get();b=get();c=get();

        add(a,b,c);

    }

    vis[1]=1;

    dfs(1);

    for(int i=1;i<=n;i++)

      for(int j=1;j<=n;j++)

        dp[i][j]=inf;

    solve(1);

    printf("%lld",dp[1][k]);

    return 0;

}