bzoj4033: [HAOI2015]树上染色

#include<cstdio>
#include<cstring> 
#include<algorithm> 
  
#define LL long long 
int n,K; 
const int maxn = 2007; 
struct node { 
    int v,w,next; 
} edge[maxn << 1]; 
int num = 0,head[maxn],size[maxn];  
inline void add_edge(int u,int v,int w) { 
    edge[++ num].v = v,edge[num].next = head[u];edge[num].w = w; head[u] = num; 
}  
LL dp[maxn][maxn];  
void dfs(int x,int fa) {  
    size[x] = 1;dp[x][0] = dp[x][1] = 0; 
    for(int i = head[x];i;i = edge[i].next) { 
        int v = edge[i].v; 
        if(v == fa) continue;  dfs(v,x); size[x] += size[v]; 
    } 
    for(int i = head[x];i;i = edge[i].next) { 
    int v = edge[i].v; 
        if(v == fa)continue; 
        for(int j = std::min(size[x],K);j >= 0;-- j) { 
        for(int k = 0;k <= std::min(j,size[v]);++ k) {
                    if(dp[x][j - k] >= 0) {
                    LL tmp = (LL)k * (K - k) * edge[i].w + (LL)(n - K - (size[v] - k)) * (size[v] - k) * edge[i].w;         
                    dp[x][j] = std::max(dp[x][j - k] + dp[v][k] + tmp,dp[x][j]);   
            }    
        } 
        } 
    } 
} 
int main() { 
    memset(dp,-1,sizeof dp); 
    scanf("%d%d",&n,&K); 
    for(int u,v,w,i = 1;i < n;++ i) { 
        scanf("%d%d%d",&u,&v,&w);
        add_edge(u,v,w);
        add_edge(v,u,w); 
    } 
    dfs(1,0); 
    printf("%lld\n",dp[1][K]); 
    return 0; 
}