牛客-HAOI2015]树上染色(树形dp)

题目链接
题面：

题解：
设$sz[i]$ 为以$i$为根的子树的大小，设$dp[i][j]$ 为以$i$ 为根的子树，选$j$ 个点染成黑色的最大收益。

接下来就是考虑怎么去计算出收益，因为不能暴力枚举黑点算距离，所以我们可以考虑每条边的贡献，

贡献=(子树内黑点 $\cdot$ 子树外黑点+子树内白点 $\cdot$ 子树外白点) $\cdot$ 边权

转移方程就是类似于背包，枚举当前黑点数和子树的黑点数

代码：

#include<cstdio>
#include<iostream>
#include<algorithm>
#include<cstring>
#include<cmath>
#include<queue>
#include<map>
#include<stack>
#include<set>
#include<ctime>
#define iss ios::sync_with_stdio(false)
using namespace std;
typedef unsigned long long ull;
typedef long long ll;
const int mod=1e9+7;
const int MAXN=2e3+5;
const ll inf=0x3f3f3f3f3f;
struct node
{
    int to;
    int next;
    ll cost;
}e[MAXN<<1];
int cnt,n,k;
int head[MAXN];
int sz[MAXN];
ll dp[MAXN][MAXN];
void add(int u,int v,int w){
    e[cnt].to=v;
    e[cnt].next=head[u];
    e[cnt].cost=w;
    head[u]=cnt++;
}
void dfs(int u,int f)
{
    dp[u][0]=dp[u][1]=0;
    sz[u]=1;
    for(int i=head[u];i!=-1;i=e[i].next)
    {
        int v=e[i].to;
        if(v==f) continue;
        dfs(v,u);
        sz[u]+=sz[v];
        for(int j=min(k,sz[u]);j>=0;j--)
        {
            for(int z=0;z<=min(j,sz[v]);z--)
            {
                ll cnt=z*(k-z)+(sz[v]-z)*(n-k-sz[v]+z);
                dp[u][j]=max(dp[u][j],dp[u][j-z]+dp[v][z]+cnt*e[i].cost);
            }
        }
    }
}
int main()
{
    memset(dp,-inf,sizeof dp);
    memset(head,-1,sizeof  head);
    scanf("%d%d",&n,&k);
    for(int i=1;i<=n-1;i++){
        int u,v;ll w;
        scanf("%d%d%lld",&u,&v,&w);
        add(u,v,w);
        add(v,u,w);
    }
    dfs(1,-1);
    printf("%lld\n",dp[1][k]);
}