Educational Codeforces Round 2 E. Lomsat gelral(樹上啓發式合併dsu on tree)

題目鏈接：https://codeforces.com/contest/600/problem/E

 

題目大意：

  求每一棵子樹各自的出現最多次的顏色編號的和,$n≤1e5$

 

題目思路：

  解決的最簡單方法自然是直接暴力，但是可以發現對於1e5的範圍無能爲力。所以這裏要引入一種新的算法，針對不帶修改的子樹問題有着獨特的優勢。
  dsu是並查集的意思，但是這個算法跟並查集沒有一點關係，唯一有點搭邊的是並查集的一種優化方法，就是每次都讓小的部分併入大的來縮短時間。
  首先照搬樹鏈剖分的概念，重兒子就是最大子樹的根節點，dsu的思想就是，先暴力算所有輕兒子的情況，每個輕兒子算完都得清空自己的貢獻，因爲前一個孩子的計數會影響後面的孩子，算完後最後計算重孩子的情況。因爲重孩子已經是最後一個處理的孩子了，剩下來的就是整個子樹的大家族，也就是自己本身，加上重孩子和所有輕孩子的情況，所以重孩子是需要的一部分，不需要再刪掉，然後暴力再把輕孩子和自己算上即可。時間複雜度的節省就在少算了一次重孩子上，根據大佬的複雜度推算，是$O(nlogn)$的。
  代碼中cnt記錄當前統計情況，dfs1處理每個點的重兒子，dfs2正式計算。flag是指是否刪除當前貢獻，1是刪除，0不刪除。首先先計算輕兒子，所以遇到重兒子跳過。這時候算的輕兒子需要去掉貢獻。然後如果有重兒子的話就處理重兒子，重兒子的情況保留，同時更新Son，跟後面用來更新的add說一聲現在的重兒子是Son，遇到它不用算了已經算過了。然後後面add的最後一個參數就是val，就是加或者刪。加完孩子後就得到了當前點代表的子樹的情況。Son變成0，因爲如果要刪當前的情況的話，得無差別全部刪光，所以得把Son變成不會遍歷到的點，才能保證全部刪掉，否則重兒子的情況就被保留了。

 

以下是代碼：

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

#define ll long long
#define rep(i,a,b) for(ll i=a;i<=b;i++)
#define per(i,a,b) for(ll i=a;i>=b;i--)
const ll MAXN =2e5+5;
const ll MAXM = 4e7+5;
const ll MOD = 998244353;
int n,c[MAXN],siz[MAXN],son[MAXN],cnt[MAXN],x,y;
ll ans[MAXN];
vector<int>v[MAXN];
void dfs1(int u,int fa){
    siz[u]=1,son[u]=0;
    int len=v[u].size();
    rep(i,0,len-1){
        int y=v[u][i];
        if(y==fa)continue;
        dfs1(y,u);
        siz[u]+=siz[y];
        if(siz[y]>siz[son[u]]){
            son[u]=y;
        }
    }
}
int Son,maxx;
ll sum;
void add(int u,int fa,int val){
    cnt[c[u]]+=val;
    if(cnt[c[u]]>maxx){
        maxx=cnt[c[u]];
        sum=c[u];
    }
    else if(cnt[c[u]]==maxx){
        sum+=c[u];
    }
    int len=v[u].size();
    rep(i,0,len-1){
        int y=v[u][i];
        if(y==fa||y==Son)continue;
        add(y,u,val);
    }
}
void dfs2(int u,int fa,int flag){
    int len=v[u].size();
    rep(i,0,len-1){
        int y=v[u][i];
        if(y==fa||y==son[u])continue;
        dfs2(y,u,1);
    }
    if(son[u])dfs2(son[u],u,0),Son=son[u];
    add(u,fa,1);
    Son=0;
    ans[u]=sum;
    if(flag)add(u,fa,-1),sum=maxx=0;
}
int main()
{
    while(cin>>n){
        memset(cnt,0,sizeof(cnt));
        Son=maxx=sum=0;
        rep(i,1,n)cin>>c[i];
        rep(i,1,n)v[i].clear();
        rep(i,1,n-1){
            cin>>x>>y;
            v[x].push_back(y);
            v[y].push_back(x);
        }
        dfs1(1,0);
        dfs2(1,0,0);
        rep(i,1,n){
            cout<<ans[i]<<" ";
        }cout<<endl;
    }
}