2017 Multi-University Training Contest - Team 3 :RXD and dividing

题目链接

#include <iostream>
#include <stdio.h>
#include <string.h>
#include <algorithm>

using namespace std;

const int maxn = 1e7;
int treeSize[maxn];
int head[maxn];
struct Edge
{
    int to;
    int weight;
    int next;
}edge[maxn];
void addEdge(int n)
{
    int a,b,c;
    memset(head,-1,sizeof(head));
    for(int i = 0; i < n-1; i++)
    {
        scanf("%d%d%d",&a,&b,&c);
        edge[i].to = b;
        edge[i].weight = c;
        edge[i].next = head[a];
        head[a] = i;
    }
}
void dfs(int u)
{
    treeSize[u] = 1;
    for(int i = head[u]; i != -1; i = edge[i].next)
    {
        int v = edge[i].to;
        dfs(v);
        treeSize[u] += treeSize[v];
    }
}
long long calAns(int n,int k)
{
    long long ans = 0;
    for(int i = 0; i < n-1; i++)
    {
        int v = edge[i].to;
        ans = ans + (long long)edge[i].weight*min(treeSize[v],k);
    }
    return ans;
}
int main()
{
    int n,k;
    while(~scanf("%d%d",&n,&k))
    {
        addEdge(n);
        dfs(1);
        printf("%lld\n",calAns(n,k));
    }
    return 0;
}