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;
}