线段树合并
对每个点开一个权值线段树
插在颜色上,维护max和sum这两个变量
不要直接#define int long long
要把需要的才开long long
不然还是会MLE
#include<bits/stdc++.h>
// #define int long long
using namespace std;
const int N=1e5+10,M=2*N;
int h[N],e[M],ne[M],idx;
int n,m,c[N],root[N],cnt;
long long ans[N];
struct node{
int l,r;
long long sum;//主要颜色的编号和
int max;//出现最多颜色的次数
#define ls tr[u].l
#define rs tr[u].r
}tr[N*100];
void add(int a,int b){
e[++idx]=b,ne[idx]=h[a],h[a]=idx;
}
void pushup(int u){
tr[u].max=max(tr[ls].max,tr[rs].max);
if(tr[ls].max==tr[rs].max)tr[u].sum=tr[ls].sum+tr[rs].sum;
if(tr[ls].max >tr[rs].max)tr[u].sum=tr[ls].sum;
if(tr[ls].max <tr[rs].max)tr[u].sum=tr[rs].sum;
}
void pushdown(int u){
}
void modify(int &o,int l,int r,int i,int k){
if(!o)o=++cnt;
if(l==r){
tr[o].max=1;
tr[o].sum=k;
return;
}
int mid=l+r>>1;
if(i<=mid)modify(tr[o].l,l,mid,i,k);
else modify(tr[o].r,mid+1,r,i,k);
pushup(o);
}
int merge(int a,int b,int l,int r){//线段树合并
if(!a)return b;
if(!b)return a;//哪个没有就返回另一个点的编号
if(l==r){
tr[a].max+=tr[b].max;
return a;
}
int mid=l+r>>1;
tr[a].l=merge(tr[a].l,tr[b].l,l,mid);//合并a,b的左儿子
tr[a].r=merge(tr[a].r,tr[b].r,mid+1,r);//合并a,b的右儿子
pushup(a);
return a;
}
void cal(int u,int fa=0){
for(int i=h[u];i;i=ne[i]){
int j=e[i];
if(j==fa)continue;
cal(j,u);
root[u]=merge(root[u],root[j],1,n);
}
ans[u]=tr[root[u]].sum;
}
signed main(){
ios::sync_with_stdio(0);cin.tie(0);cout.tie(0);
cin>>n;
for(int i=1;i<=n;i++)cin>>c[i];
for(int i=1;i<n;i++){
int a,b;
cin>>a>>b;
add(a,b);
add(b,a);
}
for(int i=1;i<=n;i++)modify(root[i],1,n,c[i],c[i]);
cal(1);
for(int i=1;i<=n;i++)cout<<ans[i]<<" ";
}