codeforces1223E Paint the Tree

https://codeforces.com/problemset/problem/1223/E

还好unrated ，每次打div1都像加了降智buff，各种水题写不出。

定一个dpk表示他向儿子最多选k条边的子树最大值，dpk1表示最多选k-1条边的子树最大值。

那么对于一个子树的根节点，要么就选e[i].w+dpk1[v],要么就选dpk[v]，而e[i].w+dpk1[v]的选择是有个数限制的，所以我们直接把e[i].w+dpk1[v]-dpk[v]丢进数组，然后贪心取前k个大和前k-1个大且>0的，然后再把所有子节点v的dpk[v]加上，就能得到dpk[u]和dpk1[u]了。

#include<bits/stdc++.h>
#define maxl 500010
#define mkp make_pair
using namespace std;

int n,k,cnt;
int ehead[maxl];
long long dpk[maxl],dpk1[maxl];
struct ed
{
	int to,nxt,w;
}e[maxl<<1];
typedef pair<long long,int> p;
typedef pair<p,int> pp;
vector <p> a[maxl]; 
bool in[maxl];

inline void add(int u,int v,int w)
{
	e[++cnt].to=v;e[cnt].w=w;
	e[cnt].nxt=ehead[u];ehead[u]=cnt;
}

inline void prework()
{
	scanf("%d%d",&n,&k);
	for(int i=1;i<=n;i++)
	{
		ehead[i]=0;
		a[i].clear();
		in[i]=false;
		dpk[i]=dpk1[i]=0;
	}
	int u,v,w;cnt=1;
	for(int i=1;i<=n-1;i++)
	{
		scanf("%d%d%d",&u,&v,&w);
		add(u,v,w);add(v,u,w);
	}
}

inline bool cmp(const p &a,const p &b)
{
	return a.first>b.first;
}

inline void dfs(int u,int fa)
{
	int v;
	for(int i=ehead[u];i;i=e[i].nxt)
	{
		v=e[i].to;
		if(v==fa)
			continue;
		dfs(v,u);
		a[u].push_back(mkp(e[i].w+dpk1[v]-dpk[v],v));
	}
	int cnt=0,l=a[u].size();
	sort(a[u].begin(),a[u].end(),cmp);
	for(int i=0;i<l;i++)
	{
		v=a[u][i].second;
		if(a[u][i].first<=0)
			break;
		cnt++;
		if(cnt<=k-1)
			dpk1[u]+=a[u][i].first;
		if(cnt<=k)
			dpk[u]+=a[u][i].first;
	}
	for(int i=ehead[u];i;i=e[i].nxt)
	{
		v=e[i].to;
		if(v==fa)
			continue;
		dpk1[u]+=dpk[v];
		dpk[u]+=dpk[v];
	}
}

inline void mainwork()
{
	dfs(1,0);
}

inline void print()
{
	printf("%lld\n",dpk[1]);
}

int main()
{
	int t;
	scanf("%d",&t);
	for(int i=1;i<=t;i++)
	{
		prework();
		mainwork();
		print();
	}
	return 0;
}