树上启发式合并(dsu on tree)

CF600E

#include<bits/stdc++.h>
using namespace std;
using ll = long long;
using ull = unsigned long long;
using i128=__int128;
using pii = pair<int,int>;
using pll = pair<ll,ll>;
using db = double;
using ldb = long double;
#define F first
#define S second
#define int ll
#define debug(x) cerr<<#x<<":"<<x<<"\n";
#define debugv(v) cerr<<#v<<":\n";for(int i=0;i<v.size();i++) cerr<<v[i]<<" ";cerr<<"\n";
void solve(){
    int n;cin>>n;
    vector<int> a(n+1);
    for(int i=1;i<=n;i++){
        cin>>a[i];
    }
    vector<vector<int>> adj(n+1);
    for(int i=1;i<n;i++){
        int u,v;cin>>u>>v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    vector<int> ans(n+1),sz(n+1),deep(n+1),f(n+1),son(n+1),t(n+1);
    auto dfs1=[&](auto&& self,int u,int fa)->void {
        sz[u]=1,deep[u]=deep[fa]+1,f[u]=fa;
        for(auto v:adj[u]){
            if(v==fa) continue;
            self(self,v,u);
            if(sz[son[u]]<sz[v]) son[u]=v;
            sz[u]+=sz[v];
        }
    };
    dfs1(dfs1,1,0);
    vector<int> cnt(n+1);
    int mx=0,sum=0;
    auto add=[&](auto&& self,int u,int sn)->void {
        cnt[a[u]]++;
        if(cnt[a[u]]>mx) mx=cnt[a[u]],sum=a[u];
        else if(cnt[a[u]]==mx) sum+=a[u];
        for(auto v:adj[u]){
            if(v==f[u]||v==sn) continue;
            self(self,v,sn);
        }
    };
    auto sub=[&](auto&& self,int u)->void {
        cnt[a[u]]--;
        for(auto v:adj[u]){
            if(v==f[u]) continue;
            self(self,v);
        }
    };
    auto dfs2=[&](auto&& self,int u,int fa,int op)->void {
        for(auto v:adj[u]){
            if(v==son[u]||v==fa) continue;
            self(self,v,u,0);
        }
        if(son[u]) self(self,son[u],u,1);
        add(add,u,son[u]);
        ans[u]=sum;
        if(!op){
            sub(sub,u);
            mx=sum=0;
        }
    };
    dfs2(dfs2,1,0,1);
    for(int i=1;i<=n;i++){
        cout<<ans[i]<<" ";
    }
}
signed main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr),cout.tie(nullptr);
    int t=1;
    // cin>>t;
    while(t--) solve();
    return 0;
}