HDU6060 RXD and dividing【DFS】

题意：一颗树，把2~n点分成k部分，求1到k部分连通的最大权重

思路：考虑每条边的贡献，u -> v（v连向1），这条边将被用min(sz(u),k)，sz是以u为根子树的点的数量

#include<stdio.h>
#include<iostream>
#include<string.h>
#include<string>
#include<stdlib.h>
#include<math.h>
#include<vector>
#include<list>
#include<map>
#include<stack>
#include<queue>
#include<algorithm>
#include<numeric>
#include<functional>
using namespace std;
typedef long long ll;
const int maxn = 1e6+5;
int sz[maxn],n,k,vis[maxn];
vector<int> v[maxn],dis[maxn];
long long ans;

void init(int x)
{
	for(int i = 0 ; i <= x; i++)
	{
		v[i].clear();
		dis[i].clear();
		sz[i] = 1;
	}
	memset(vis,0,sizeof vis);
}
void add(int a,int b,int c)
{
	v[a].push_back(b);
	v[b].push_back(a);
	dis[a].push_back(c);
	dis[b].push_back(c);
}

int dfs(int x)
{
	vis[x] = 1;
	for(int i = 0; i < v[x].size(); i++)
	{
		int temp = v[x][i];
		if(vis[temp]) continue;
		sz[x] += dfs(temp);
		ans += (long long)dis[x][i] * min(sz[temp],k);
	}
	return sz[x];
}

int main(void)
{
	while(scanf("%d%d",&n,&k)!=EOF)
	{
		init(n);
		for(int i = 1; i < n; i++)
		{
			int a,b,c;
			scanf("%d%d%d",&a,&b,&c);
			add(a,b,c);
		}
		ans = 0ll;
		dfs(1);
		printf("%lld\n",ans);
	}
	return 0;
}