树上启发式合并 DSU on Tree

更新日志

2025/01/07：开工。

---

概念

树上启发式合并，可以一定程度上减小合并操作的复杂度，或者保证正确性。

思路

对于每一个节点，我们都找出它的最重儿子，也就是子节点个数最多的儿子。如有多个，任选一个。

首先统计其他轻儿子的答案（如果无需统计每个节点的答案，就不用了。）。

下面正式开始启发式合并。

1. 跑一遍重儿子，获取答案。
2. 直接把根节点答案合并进去，作为根节点的答案。
3. 对于所有轻儿子，再次搜索一遍，更新根节点的答案。

这样复杂度是 \(n\log n\) 的。

例题

CF600E

#include<bits/stdc++.h>
using namespace std;

typedef long long ll;
typedef unsigned long long ull;
typedef __int128 i128;
typedef double db;
typedef long double ld;
typedef pair<int,int> pii;
typedef pair<ll,ll> pll;
typedef pair<int,ll> pil;
typedef pair<ll,int> pli;
template <typename Type>
using vec=vector<Type>;
template <typename Type>
using grheap=priority_queue<Type>;
template <typename Type>
using lrheap=priority_queue<Type,vector<Type>,greater<Type> >;
#define fir first
#define sec second
#define pub push_back
#define pob pop_back
#define puf push_front
#define pof pop_front
#define chmax(a,b) a=max(a,b)
#define chmin(a,b) a=min(a,b)
#define rep(i,x,y) for(int i=x;i<=y;i++)
#define per(i,x,y) for(int i=x;i>=y;i--)

const int inf=0x3f3f3f3f;
const ll INF=0x3f3f3f3f3f3f3f3f;
const int mod=998244353;

const int N=1e5+5;

int n;
int c[N];

vec<int> vs[N];

map<int,int> mp;

ll ans[N];
int mx;
ll sum;

int sz[N],ms[N];
void init(int now,int fid){
    sz[now]=1;
    for(auto nxt:vs[now]){
        if(nxt==fid)continue;
        init(nxt,now);
        sz[now]+=sz[nxt];
        if(sz[nxt]>=sz[ms[now]])ms[now]=nxt;
    }
}

void dfs1(int now,int fid){
    mp[c[now]]++;
    if(mp[c[now]]>mx)mx=mp[c[now]],sum=c[now];
    else if(mp[c[now]]==mx)sum+=c[now];
    for(auto nxt:vs[now]){
        if(nxt==fid)continue;
        dfs1(nxt,now);
    }
}

void dfs(int now,int fid){
    for(auto nxt:vs[now]){
        if(nxt==fid||nxt==ms[now])continue;
        dfs(nxt,now);
        mp.clear();
        mx=0;sum=0;
    }
    if(ms[now])dfs(ms[now],now);
    mp[c[now]]++;
    if(mp[c[now]]>mx)mx=mp[c[now]],sum=c[now];
    else if(mp[c[now]]==mx)sum+=c[now];
    for(auto nxt:vs[now]){
        if(nxt==fid||nxt==ms[now])continue;
        dfs1(nxt,now);
    }
    ans[now]=sum;
}

int main(){
    ios::sync_with_stdio(false);
    cin.tie(0);cout.tie(0);
    cin>>n;
    rep(i,1,n)cin>>c[i];
    rep(i,2,n){
        int u,v;cin>>u>>v;
        vs[u].pub(v);
        vs[v].pub(u);
    }
    init(1,1);
    dfs(1,1);
    rep(i,1,n)cout<<ans[i]<<" ";
    return 0;
}