bzoj 4033: [HAOI2015]树上染色 [树形DP]

4033: [HAOI2015]树上染色

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <cmath>
using namespace std;
typedef long long ll;
const int N = 2005, P = 1e9+7;
inline int read() {
char c=getchar(); int x=0,f=1;
while(c<'0'||c>'9') {if(c=='-')f=-1;c=getchar();}
while(c>='0'&&c<='9') {x=x*10+c-'0';c=getchar();}
return x*f;
}

int n, m, mm, u, v;
struct edge{int v, ne, w;} e[N<<1];
int cnt, h
;
inline void ins(int u, int v, int w) {
e[++cnt] = (edge){v, h[u], w}; h[u] = cnt;
e[++cnt] = (edge){u, h[v], w}; h[v] = cnt;
}

ll f

; int size
;
void dp(int u, int fa) { //printf("dp %d %d\n", u, fa);
size[u] = 1;
for(int i=h[u]; i; i=e[i].ne) {
int v = e[i].v, w = e[i].w;
if(v == fa) continue;
dp(v, u);
for(int j = min(size[u] + size[v], m); j >= 0; j--) {
int _ = min(j, size[v]);
for(int k = max(0, j - size[u]); k <= _; k++)
f[u][j] = max(f[u][j], f[v][k] + f[u][j-k] + (ll) w * ( k * (m-k) + (size[v] - k) * (mm - size[v] + k) ) );
}
size[u] += size[v];
}
//printf("look %d  %d\n", u, size[u]);
//for(int i=0; i<=min(size[u], m); i++) printf("f %d %d  %lld\n", u, i, f[u][i]);
//puts("end\n");
}
int main() {
//freopen("in", "r", stdin);
freopen("haoi2015_t1.in", "r", stdin);
freopen("haoi2015_t1.out", "w", stdout);
n = read(); m = read(); mm = n - m;
for(int i=1; i<n; i++) u = read(), v = read(), ins(u, v, read());
dp(1, 0);
printf("%lld\n", f[1][m]);
}