BZOJ4033 [HAOI2015]樹上染色

本來是考慮， $ f[x][i][0/1] $ 表示 $ x $ 子樹中有$i$個黑點，且 $ x $ 是白點/黑點。但是這裡的答案是要統計不同的子樹的貢獻的。所以就gg了。
看了題解。
應該是要設$f[x][i]$表示$x$子樹中有$i$個黑點，對答案的貢獻。
轉移的時候，就可以單獨計算出$x->y$(y是x的兒子)這條邊的貢獻。
貢獻怎麼算呢？就是統計一下$y$內有多少黑(白)點、$y$外有多少黑(白)點，算一下有多少對，最後乘上$x->y$的邊權。

#include<iostream>
#include<cstdio>
#include<algorithm>
#include<cstring>
using namespace std;
#define ll long long
#define rg register
#define inf 0x3f3f3f3f
const int MX=4005;
template <typename _Tp> inline void read(_Tp&x){
    char c=getchar();x=0;bool b=0;
    while(c!='-'&&!isdigit(c))c=getchar();if(c=='-'){c=getchar();b=1;}
    while(isdigit(c)){x=x*10+c-'0';c=getchar();}if(b)x=-x;return;
}
int n,k,cnt;
int head[MX],sz[MX];
ll f[MX][MX];
struct Edge{int v,nxt,w;}edge[MX<<1];
inline void add_edge(int x,int y,int z){edge[++cnt].v=y;edge[cnt].w=z;edge[cnt].nxt=head[x];head[x]=cnt;}
void dfs(int x,int fa){
    sz[x]=1;f[x][0]=f[x][1]=0;
    for(int i=head[x];i;i=edge[i].nxt){
        int y=edge[i].v;
        if(y==fa) continue;
        dfs(y,x);
        for(int t=min(k,sz[x]+sz[y]);~t;t--)
            for(int j=max(0,t-sz[x]);j<=min(sz[y],t);j++){//枚舉子樹y中黑點的數量
                ll w=1LL*(1LL*j*(k-j)+1LL*(sz[y]-j)*(n-k-sz[y]+j))*edge[i].w;//當前邊的貢獻
                f[x][t]=max(f[x][t],f[x][t-j]+f[y][j]+w);
            }
        sz[x]+=sz[y];
    }
}
int main(){
    // freopen("a.in","r",stdin);
    read(n);read(k);
    int x,y,z;
    for(int i=1;i<n;i++){
        read(x);read(y);read(z);
        add_edge(x,y,z),add_edge(y,x,z);
    }
    memset(f,-inf,sizeof(f));
    dfs(1,-1);

    printf("%lld\n",f[1][k]);
    return 0;
}