树上启发式合并学习笔记

洛谷日报

例题

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
#define inf 0x3f3f3f3f
#define endl '\n'
struct edge{
    int nxt,to;
}e[200050];
ll col[100050],cnt[100050],ans[100050],sum,mx;
int head[100050],a[100050],ct,sz[100050],son[100050];
void add(int u,int v)
{
    e[++ct].to=v;
    e[ct].nxt=head[u];
    head[u]=ct;
}
void dfs1(int u,int fa) //求重儿子
{
    sz[u]=1;
    for(int i=head[u];i;i=e[i].nxt){
        int v=e[i].to;
        if(v!=fa){
            dfs1(v,u); sz[u]+=sz[v];
            if(sz[v]>sz[son[u]]){
                son[u]=v;
            }
        }
    }
}
void change(int u,int fa,int val,int p)
{
    cnt[col[u]]+=val;
    if(cnt[col[u]]>mx){
        mx=cnt[col[u]]; sum=col[u];
    }
    else if(cnt[col[u]]==mx){
        sum+=col[u];
    }
    for(int i=head[u];i;i=e[i].nxt){
        int v=e[i].to;
        if(v==fa||v==p){
            continue;
        }
        change(v,u,val,p);
    }
}
void dfs2(int u,int fa,int op)
{
    for(int i=head[u];i;i=e[i].nxt){
        int v=e[i].to;
        if(v==fa||v==son[u]) continue;
        dfs2(v,u,0);
    }
    if(son[u]) dfs2(son[u],u,1); //如果是重儿子的话
    change(u,fa,1,son[u]); 
    ans[u]=sum;
    if(op==0){ //如果是轻儿子的话
        change(u,fa,-1,0); //重置清零
        sum=mx=0;
    }
}
int main(){
ios::sync_with_stdio(0),cin.tie(0),cout.tie(0);
int n;cin>>n;
for(int i=1;i<=n;i++){
    cin>>col[i];
}
for(int i=1;i<n;i++){
    int x,y; cin>>x>>y;
    add(x,y); add(y,x);
}
dfs1(1,0); dfs2(1,0,0);
for(int i=1;i<=n;i++){
    cout<<ans[i]<<" ";
}
    return 0;
}