HDU 6060 RXD and dividing

题意：
给出一个具有n个顶点的树， 将其中的2~n分成k部分， 求1号顶点到分成k块后各个顶点的最大权值和。
思路：
将2~n的结点编号{1~k}，每条边的权乘以其下子树的大小与k的较小值。
代码：

# include <bits/stdc++.h>
using namespace std;
typedef long long LL;

const int INF = 0x3f3f3f3f;
const int Q = 1e9+7;
const int N = 1e6+10;

int n, k, tot;
int head[N], f[N], si[N], w[N];
struct Edge {
    int fr, to, co, ne;
} E[N * 2];

void add (int u, int v, int w)
{
    E[tot].fr = u;
    E[tot].to = v;
    E[tot].co = w;
    E[tot].ne = head[u];
    head[u] = tot++;
}
void dfs (int u, int fa)
{
    si[u] = 1;
    for (int i=head[u]; i!=-1; i=E[i].ne) {
        int to = E[i].to;
        if (to == fa) continue;
        w[to] = E[i].co;
        dfs(to, u);
        si[u] += si[to];
    }
}
void solve ()
{
    LL ans = 0;
    tot = 0;
    memset (head, -1, sizeof (head));
    memset (si, 0, sizeof (si));
    for (int i=0; i<n-1; ++i) {
        int u, v, w;
        scanf ("%d%d%d", &u, &v, &w);
        add(u, v, w);
        add(v, u, w);
    }
    dfs (1, -1);
    for (int i=2; i<=n; ++i) {
        ans += (LL)w[i] * min(si[i], k);
    }
    printf ("%lld\n", ans);
}

int main ()
{
    while (scanf ("%d%d", &n, &k) != EOF) {
        solve ();
    }

    return 0;
}