洛谷P3177 [HAOI2015]树上染色

洛谷P3177

今天刷知乎，他给我推荐了一个ACM竞赛好题，然后那个作者强推我们做一个树上dp必做题，然后中午吃饭的时候我就好奇看了看是什么题，中午吃饭的时候想了一下思路，但是感觉做不了，于是到晚上找出来开始做。
题意： 给你一棵带权，然后染黑白两色，要求染k个黑色，然后我们得分为黑点之间两两距离之和，和白点之间两两距离之和。
思路：
开始思路非常明确，就是$dp[i][j]$表示i为树根的子树中有j个黑色的节点的答案数，然后就就想到了上司的舞会转移法，其实不能转移，状态太多了越想越乱，后来理了一下思路，就是条件（黑色和白色两两之间的距离肯定要转化一步），一看题解，果然，我们算的是每条边的贡献,然后每条边的贡献就是，边两边同色点数量的乘积，然后树上背包，大悟，整了一下，但是没注意初始化，dp数组，一定要注意状态转移的，有可能有转移不到的，然后就是就需要特判。
然后看了题解发现还要初始化$dp[u][0]$和$dp[u][1]$因为这两个是必达状态

其实是一道非常好的树上背包题目
AC代码：

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int maxn = 2005;
ll dp[maxn][maxn];
struct edge{
    int v,next;
    ll w;
}e[maxn<<1];
int sz[maxn];
int head[maxn],cnt = 0;
void add(int u,int v,int w){
    e[cnt].v = v;
    e[cnt].w = w;
    e[cnt].next = head[u];
    head[u] = cnt++;
}
int n,k;
void dfs(int u,int fa){
    sz[u] = 1;
    memset(dp[u],-1,sizeof dp[u]);
    dp[u][0] = 0;dp[u][1] = 0;
    for(int i = head[u];~i;i=e[i].next){
        int v = e[i].v;
        if(v==fa)continue;
        dfs(v,u);
        sz[u] += sz[v];
    }
    for(int i = head[u];~i;i=e[i].next){
        int v = e[i].v;
        int w = e[i].w;
        if(v==fa)continue;
        for(int t = min(k,sz[u]);t >= 0; t--){
            for(int j = 0;j <= min(t,sz[v]);j++){
                if(dp[u][t-j]==-1)continue;
                ll tot = 1LL*j*(k-j) + (sz[v] - j) * 1LL*(n - sz[v] - k + j);
                tot *= w;
                dp[u][t] = max(dp[u][t],dp[u][t-j] + dp[v][j] + tot);
            }
        }
    }
}
int main(){
    scanf("%d%d",&n,&k);
    memset(head,-1,sizeof head);
    for(int i = 0 ;i < n - 1;i ++){
        int u,v,w;
        scanf("%d%d%d",&u,&v,&w);
        add(u,v,w);
        add(v,u,w);
    }
    dfs(1,-1);
    printf("%lld",dp[1][k]);
}