cf 600E. Lomsat gelral 树上DSU

Sample input
15
1 2 3 1 2 3 3 1 1 3 2 2 1 2 3
1 2
1 3
1 4
1 14
1 15
2 5
2 6
2 7
3 8
3 9
3 10
4 11
4 12
4 13

Sample output
6 5 4 3 2 3 3 1 1 3 2 2 1 2 3

最近听群里大佬们讨论,于是也想学一学
主要就是树上的启发式合并
统计子树信息的时候,可以把轻的子树的信息合并到重的子树上去
由于一次合并起码会使得被合并的size增加一倍轻的数量
所以每个点最多被合并logn次

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

#define REP(i,j,k) for(int i=j; i<k; i++)

int n;
constexpr int N = 1e5+10;
int color[N];
vector<int> g[N];
unordered_map<int, int> cnt[N]; //计算每个点的子树的颜色的个数
long long ans[N]; //记录每个点的儿子
int sze[N];//记录儿子的个数，这样来启发的合并
int dsu[N]; //记录合并到哪个节点了
int mx[N];
void add(int des, int c, int v){
    if(mx[des]<c) {  //最大值变化了需要被更新
        mx[des] = c;
        ans[des] = v;
    }
    else if(mx[des]==c) ans[des]+=v;
}
void merge(int v, int des, int s){
    for(auto&e:cnt[s]){  //将另一个子树的信息合并过来
        add(v, cnt[des][e.first]+=e.second, e.first);
    }
}
void dfs(int cur, int fa){
    sze[cur] = 1;
    if(g[cur].size()==0||g[cur].size()==1&&g[cur].back()==fa){//没有儿子节点
        ans[cur] = color[cur];
        dsu[cur] = cur;  //直接指定使用当前的这个点对应的map，这样不会冲突
        cnt[dsu[cur]][color[cur]] = 1;
        mx[cur] = 1;
        return;
    } 
    int big=N-1;
    for(auto&e:g[cur]){
        if(e==fa) continue;  
        dfs(e, cur); //对儿子进行搜索
        sze[cur] += sze[e];
        if(sze[big]<sze[e]) big=e; //记录最大的儿子
    }

    dsu[cur] = dsu[big];  //因为儿子的以及统计过了，那么破坏他的也没事
    ans[cur] = ans[big];  //先更新
    mx[cur] = mx[big];  //暂时存这个值

    cnt[dsu[cur]][color[cur]] += 1;  //改变这个颜色的数量
    add(cur, cnt[dsu[cur]][color[cur]],color[cur]);  


    //合并size更小的子树的节点到当前点
    for(auto&e:g[cur]){
        if(e!=big){ //合并非最大的点
            merge(cur, dsu[cur], dsu[e]);
        }
    }
}

int main(){
    // freopen("in.dat", "r", stdin);
    cin>>n;
    REP(i,0,n) cin>>color[i+1];
    int from, to;
    // int root=0;
    REP(i,0,n-1){
        cin>>from>>to;
        g[from].push_back(to);
        g[to].push_back(from);
    }
    dfs(1,-1);
    for_each(ans+1, ans+n+1, [](long long&x){
        cout<<x<<" ";
    });
    return 0;
}