山东理工大学第十二届ACM程序设计竞赛 - Cut the tree(树上启发式合并+线段树)

题目链接：点击查看

题目大意：给一个具有 N 个节点的有根树，以 1 号节点为根，节点编号从 1 开始，点有点权。树的第 H 层权值为深度为 H 的所有点的点权之和。树的总权值为所有层权值的最大值。问分别割掉以 1,2,..,N 为根的子树后，剩余树的总权值为多少。

题目分析：因为题目提示了，弱化数据后支持双 log 的算法。。所以不难想到树上启发式合并+线段树，两个模板套起来就过了，线段树负责处理区间最大值问题，每次更新子树信息的时候，单点更新一下相关层的最大值就好了，记得回溯

代码：
 

#include<iostream>
#include<cstdio>
#include<string>
#include<ctime>
#include<cmath>
#include<cstring>
#include<algorithm>
#include<stack>
#include<climits>
#include<queue>
#include<map>
#include<set>
#include<sstream>
#include<cassert>
using namespace std;

typedef long long LL;

typedef unsigned long long ull;

const LL inf=0x3f3f3f3f;

const int N=1e5+100;

LL a[N],sum[N],ans[N];

struct Node
{
	int l,r;
	LL mmax;
}tree[N<<2];

void pushup(int k)
{
	tree[k].mmax=max(tree[k<<1].mmax,tree[k<<1|1].mmax);
}

void build(int k,int l,int r)
{
	tree[k].l=l;
	tree[k].r=r;
	if(l==r)
	{
		tree[k].mmax=sum[l];
		return;
	}
	int mid=l+r>>1;
	build(k<<1,l,mid);
	build(k<<1|1,mid+1,r);
	pushup(k);
}

void update(int k,int pos,int val)
{
	if(tree[k].l==tree[k].r)
	{
		tree[k].mmax+=val;
		return;
	}
	int mid=tree[k].l+tree[k].r>>1;
	if(pos<=mid)
		update(k<<1,pos,val);
	else
		update(k<<1|1,pos,val);
	pushup(k);
}

int num[N],son[N],deep[N],max_deep;

bool vis[N];
 
vector<int>node[N];

void dfs_son(int u,int fa,int dep)//树链剖分跑出重链
{
	max_deep=max(max_deep,dep);
	sum[dep]+=a[u];
	deep[u]=dep;
	son[u]=-1;
	num[u]=1;
	for(auto v:node[u])
	{
		if(v==fa)
			continue;
		dfs_son(v,u,dep+1);
		num[u]+=num[v];
		if(son[u]==-1||num[v]>num[son[u]])
			son[u]=v;
	}
}

void cal(int u,int fa,int val)//对于每个节点计算其子树的贡献
{
	update(1,deep[u],val*a[u]);
	for(auto v:node[u])
	{
		if(v==fa||vis[v])
			continue;
		cal(v,u,val);
	}
}

void dfs(int u,int fa,int keep)//启发式合并
{
	for(auto v:node[u])
	{
		if(v==fa||v==son[u])
			continue;
		dfs(v,u,0);
	}
	if(son[u]!=-1)
	{
		dfs(son[u],u,1);
		vis[son[u]]=true;
	}
	cal(u,fa,-1);
	ans[u]=tree[1].mmax;
	if(son[u]!=-1)
		vis[son[u]]=false;
	if(!keep)
		cal(u,fa,1);
}

int main()
{
#ifndef ONLINE_JUDGE
//  freopen("input.txt","r",stdin);
//  freopen("output.txt","w",stdout);
#endif
//  ios::sync_with_stdio(false);
	int n;
	scanf("%d",&n);
	for(int i=1;i<n;i++)
	{
		int u,v;
		scanf("%d%d",&u,&v);
		node[u].push_back(v);
		node[v].push_back(u);
	}
	for(int i=1;i<=n;i++)
		scanf("%lld",a+i);
	dfs_son(1,-1,1);
	build(1,1,max_deep);
	dfs(1,-1,1);
	printf("%lld",ans[1]);
	for(int i=2;i<=n;i++)
		printf(" %lld",ans[i]);
	puts("");









    return 0;
}