解题思路
\(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;
}