[HAOI2015]树上染色 [树形dp]

3177 [HAOI2015]树上染色

#include<bits/stdc++.h>
using namespace std;
#define Max(x,y) (x)>(y)?(x):(y)
#define Min(x,y) (x)>(y)?(y):(x)
#define rg register
#define ll long long
const int N=2000+5,M=0,inf=0x3f3f3f3f,P=99999997;
int n,k,size[N];
long long f[N][N];
template <class t>void rd(t &x){
    x=0;int w=0;char ch=0;
    while(!isdigit(ch)) w|=ch=='-',ch=getchar();
    while(isdigit(ch)) x=(x<<1)+(x<<3)+(ch^48),ch=getchar();
    x=w?-x:x;
}


int head[N],tot=0;
struct edge{int v,w,nxt;}e[N<<1];
void add(int u,int v,int w){
    e[++tot]=(edge){v,w,head[u]},head[u]=tot;
}

void dfs(int u,int fa){
    size[u]=1,f[u][0]=f[u][1]=0;
    for(int i=head[u];i;i=e[i].nxt)
    if(e[i].v!=fa) dfs(e[i].v,u),size[u]+=size[e[i].v];
    for(int ii=head[u],v,w;ii;ii=e[ii].nxt){
        v=e[ii].v,w=e[ii].w;
        if(v==fa) continue;
        for(int i=min(size[u],k);i>=0;--i)
        for(int j=0;j<=min(i,size[v]);++j)
        if(f[u][i-j]!=-1){
            ll val=(ll)(k-j)*j*w+(ll)(size[v]-j)*(n-size[v]-(k-j))*w;
            f[u][i]=max(f[u][i],f[u][i-j]+f[v][j]+val);
        }
    }
}

int main(){
    freopen("in.txt","r",stdin);
    rd(n),rd(k);
    for(int i=1,u,v,w;i<n;++i) rd(u),rd(v),rd(w),add(u,v,w),add(v,u,w);
    memset(f,-1,sizeof(f));
    dfs(1,-1);
    printf("%lld",f[1][k]);
    return 0;
}