CF 600E. Lomsat gelral(dsu on tree)

解题思路

　　(dsu) (on) (tree)的模板题。暴力而优雅的算法，轻儿子的信息暴力清空，重儿子的信息保留，时间复杂度(O(nlogn))

代码

#include<iostream>
#include<cstdio>
#include<cstring>
#include<cstdlib>
#include<algorithm>
#include<set>

using namespace std;
const int MAXN = 100005;
typedef long long LL;

inline int rd(){
	int x=0,f=1;char ch=getchar();
	while(!isdigit(ch)) f=ch=='-'?0:1,ch=getchar();
	while(isdigit(ch)) x=(x<<1)+(x<<3)+ch-'0',ch=getchar();
	return f?x:-x;
}

int n,c[MAXN],head[MAXN],cnt,to[MAXN<<1],nxt[MAXN<<1],num[MAXN],sum[MAXN];
int siz[MAXN],son[MAXN],fa[MAXN],Num;
LL ans[MAXN],Max[MAXN],Ans;

inline void add(int bg,int ed){
	to[++cnt]=ed,nxt[cnt]=head[bg],head[bg]=cnt;
}

inline void update(int x){
	sum[num[c[x]]]--;Max[num[c[x]]]-=c[x];
	num[c[x]]++;sum[num[c[x]]]++;Max[num[c[x]]]+=c[x];
	if(Num<=num[c[x]]) Num=num[c[x]];Ans=Max[Num];
}

inline void Clear(int x){
	sum[num[c[x]]]--;Max[num[c[x]]]-=c[x];
	num[c[x]]--;sum[num[c[x]]]++;Max[num[c[x]]]+=c[x];
	if(Num==num[c[x]]+1 && !sum[num[c[x]]+1]) Num=num[c[x]];
	Ans=Max[Num];
}

void bfs(int x){
	update(x);
	for(int i=head[x];i;i=nxt[i])
		if(to[i]!=fa[x]) bfs(to[i]);
}

void clear(int x){
	Clear(x);
	for(int i=head[x];i;i=nxt[i])
		if(to[i]!=fa[x]) clear(to[i]);
}

void dfs1(int x,int f){
	fa[x]=f;siz[x]=1;int maxson=-1,u;
	for(int i=head[x];i;i=nxt[i]){
		u=to[i];if(u==f) continue;
		dfs1(u,x);siz[x]+=siz[u];
		if(siz[u]>maxson) maxson=siz[u],son[x]=u;
	}
}

void dfs2(int x){
	for(int i=head[x];i;i=nxt[i]){
		int u=to[i];if(u==fa[x] || u==son[x]) continue;
		dfs2(u);clear(u);
	}
	if(son[x]) dfs2(son[x]);
	for(int i=head[x];i;i=nxt[i]){
		int u=to[i];if(u==fa[x] || u==son[x]) continue;
		bfs(u);
	}update(x);
	ans[x]=Ans;
}

int main(){
	n=rd();int x,y;
	for(int i=1;i<=n;i++) c[i]=rd();
	for(int i=1;i<n;i++){
		x=rd(),y=rd();
		add(x,y);add(y,x);
	}dfs1(1,0);dfs2(1);
	for(int i=1;i<=n;i++) printf("%lld ",ans[i]);
	return 0;
}