600E.Lomsat gelral (树上启发式合并)

给出一个n个点的树

求出树中每个子树出现次数最多的所有颜色的编号的和。

题解:

搜索的过程中,重儿子所在的子树搜一遍,记录颜色的出现情况

然后把轻儿子往重儿子上合并

#include<bits/stdc++.h>
using namespace std;
const int maxn=1e5+100;
int c[maxn];
vector<int> g[maxn];
long long ans[maxn];
int f[maxn];//当前节点的子树内每种颜色的出现次数
int L[maxn],R[maxn],id[maxn];
int dep[maxn];
int size[maxn];
int tot;
int son[maxn];
int Max[maxn];
int n;
void dfs1 (int x,int pre) {
	dep[x]=dep[pre]+1;
	size[x]=1;
	L[x]=++tot;
	id[tot]=x;
	for (int y:g[x]) {
		if (y==pre) continue;
		dfs1(y,x);
		size[x]+=size[y];
		if (size[son[x]]<size[y]) son[x]=y;
	}
	R[x]=tot;
} 
void cal (int x,int pre) {
	//计算x的答案
	//在cal函数里
	//重儿子的f状态已经保存  
	f[c[x]]++;
	if (f[c[x]]>Max[x]) {
		Max[x]=f[c[x]];
		ans[x]=c[x];
	}
	else if (f[c[x]]==Max[x]) {
		ans[x]+=c[x];
	}
	for (int y:g[x]) {
		//遍历除重儿子外的子树
		//把答案往重儿子上合并 
		if (y==son[x]) continue;
		if (y==pre) continue;
		for (int j=L[y];j<=R[y];j++) {
			int z=id[j];
			f[c[z]]++;
			if (f[c[z]]>Max[x]) {
				Max[x]=f[c[z]];
				ans[x]=c[z];
			}
			else if (f[c[z]]==Max[x]) {
				ans[x]+=c[z];
			}
		}
	}
//	Max[x]=0;
//	ans[x]=0;
//	for (int i=1;i<=n;i++) Max[x]=max(Max[x],f[i]);
//	for (int i=1;i<=n;i++) if (f[i]==Max[x]) ans[x]+=i; 
}
void dfs2 (int x,int pre,int kp) {
	for (int y:g[x]) {
		if (y==son[x]) continue;
		if (y==pre) continue;
		dfs2(y,x,0);
	}
	if (son[x]) {
		dfs2(son[x],x,1);
		ans[x]=ans[son[x]];
		Max[x]=Max[son[x]];
	}
	cal(x,pre);
	if (!kp) for (int i=L[x];i<=R[x];i++) f[c[id[i]]]=0;
}
int main () {
	scanf("%d",&n);
	for (int i=1;i<=n;i++) scanf("%d",c+i);
	for (int i=1;i<n;i++) {
		int x,y;
		scanf("%d%d",&x,&y);
		g[x].push_back(y);
		g[y].push_back(x);
	}
	dfs1(1,0);
	dfs2(1,0,1);
	for (int i=1;i<=n;i++) printf("%lld ",ans[i]); 
}
posted @ 2021-04-01 22:46  zlc0405  阅读(59)  评论(0编辑  收藏  举报