树上启发式合并

题目:一棵树有n个结点,每个结点都是一种颜色,每个颜色有一个编号,求树中每个子树的最多的颜色编号的和。
树上启发式合并真神奇,时间复杂度只有O(nlogn)
看下思路
第一步:我先像树剖那样,跑出重儿子。

第二步:我们dfs这个树,比如说,我们现在跑到了u节点。首先优先跑所有轻儿子,采用尾递归的方式,统计轻儿子的答案。期间维护一个通数组cnt,存的是以当前轻儿子为根的子树的颜色数量。跑完后将答案赋予这个轻儿子节点,然后清空cnt数组,开始跑下一个轻儿子。

第三步:轻儿子跑完后,我们开始跑重儿子。跑完后,注意:我们不在清空cnt数组,然后逐个再跑一遍各个轻儿子。将结果累计到cnt里面,这里就体现了我们所受的启发:小往大合并。

第四步:经第三步,我们cnt数组就存入了所有子树信息,然后结合u节点自身的信息,将这个答案赋予u就行啦。

#include<cstdio>
#define int long long
const int maxn=1e5+5;
int ver[maxn<<1],head[maxn],next[maxn<<1];
int tot=1;
void add(int x,int y){
   
	ver[++tot]=y;
	next[tot]=head[x];
	head[x]=tot;
}
int siz[maxn],son[maxn],col[maxn],cnt[maxn],ans[maxn];
void dfs_bson(int u,int f){
   
	siz[u]=1;
	for(int i=head[u];i;i=next[i]){
   
		int y=ver[i];
		if(y==f)	continue;
		dfs_bson(y,u);
		siz[u]+=siz[y];
		if(siz[y]>siz[son[u]])
			son[u]=y;
	}
}
int sum,flag,maxx;
void _count(int u,int f,int val){
   
	cnt[col[u]]+=val;
	if(cnt[col[u]]>maxx){
   
		maxx=cnt[col[u]];
		sum=col[u];
	}
	else if(cnt[col[u]]==maxx)	sum+=col[u];
	for(int i=head[u];i;i=next[i]){
   
		int y=ver[i];
		if(y==f||y==flag)	continue;
		_count(y,u,val);
	}
}
void dfs(int u,int f,bool _flag){
   
	for(int i=head[u];i;i=next[i]){
   
		int y=ver[i];
		if(y==f||y==son[u])	continue;
		dfs(y,u,0);
	}
	if(son[u]){
   
		dfs(son[u],u,1);
		flag=son[u];
	}
	_count(u,f,1);
	flag=0;
	ans[u]=sum;
	if(!_flag){
   
		_count(u,f,-1);
		sum=maxx=0;
	}
}
signed main(){
   
	int n;
	scanf("%lld",&n);
	for(int i=1;i<=n;i++)
		scanf("%lld",&col[i]);
	for(int i=1;i<n;i++){
   
		int x,y;
		scanf("%lld%lld",&x,&y);
		add(x,y);
		add(y,x);
	}
	dfs_bson(1,0);
	dfs(1,0,0);
	for(int i=1;i<=n;i++){
   
		printf("%lld",ans[i]);
		if(i!=n) printf(" ");
		
	}
	return 0;
}
全部评论

相关推荐

点赞 评论 收藏
分享
头像
11-06 10:58
已编辑
门头沟学院 嵌入式工程师
双非25想找富婆不想打工:哦,这该死的伦敦腔,我敢打赌,你简直是个天才,如果我有offer的话,我一定用offer狠狠的打在你的脸上
点赞 评论 收藏
分享
点赞 收藏 评论
分享
牛客网
牛客企业服务