Game HDU - 5242 树形dp+树链剖分

题目链接：https://vjudge.net/problem/HDU-5242

题意：n个节点的一个树，每个节点一个权值，从根节点出发，走k次，求走过权值和最大，走过的点权值变为0

题解：树链剖分的思想，我们先dfs记录每个节点权值最大的分支，在dfs一遍，有祖先的把权值都加到祖先节点上，没有祖先的先遇到的节点作为祖先，最后所有节点排个序，输出前k大即可

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N=1e5+10;
struct edge{
	int to,nex,val;
}e[N*2];
int n,k;
int head[N],len;
int val[N];
int son[N];
ll s_num[N];
ll dp[N];
void init()
{
	for(int i=0;i<=n;i++)
	{
		head[i]=-1;
		son[i]=-1;
		s_num[i]=0;
		dp[i]=0;
	}
	len=0;
}
void addedge(int x,int y,int z)
{
	e[len].to=y;
	e[len].val=z;
	e[len].nex=head[x];
	head[x]=len++;
}
void dfs1(int u,int fa)
{
	s_num[u]=val[u];
	int to;
	for(int i=head[u];i!=-1;i=e[i].nex)
	{
		to=e[i].to;
		if(to==fa) continue;
		dfs1(to,u);
		if(son[u]==-1||s_num[to]>s_num[son[u]])
		{
			son[u]=to;
		}
	}
	if(son[u]!=-1) s_num[u]+=s_num[son[u]];
}
void dfs2(int u,int fa,int rt)
{
	dp[rt]+=val[u];
	if(son[u]!=-1)
	{
		dfs2(son[u],u,rt);	
	}
	int to;
	for(int i=head[u];i!=-1;i=e[i].nex)
	{
		to=e[i].to;
		if(to==fa || to==son[u]) continue;
		dfs2(to,u,to);
	}
}
int main()
{
	int T;
	int nn=1;
	scanf("%d",&T);
	while(T--)
	{
		int x,y;
		scanf("%d%d",&n,&k);
		init();
		for(int i=1;i<=n;i++)
			scanf("%d",&val[i]);
		for(int i=1;i<n;i++)
		{
			scanf("%d%d",&x,&y);
			addedge(x,y,val[y]);
			addedge(y,x,val[x]);
		}
		dfs1(1,0);
		dfs2(1,0,1);
		sort(dp+1,dp+1+n);
		ll ans=0;
		for(int i=n;i>=n-k+1;i--)ans+=dp[i];
		printf("Case #%d: %lld\n",nn++,ans);
	}
	return 0;
}