CF 600E树上启发式合并代码入门题

// Problem: CF600E Lomsat gelral
// Contest: Luogu
// URL: https://www.luogu.com.cn/problem/CF600E
// Memory Limit: 250 MB
// Time Limit: 2000 ms
// 
// Powered by CP Editor (https://cpeditor.org)

// Problem: B. Rising Sand
// Contest: Codeforces - Codeforces Round #803 (Div. 2)
// URL: https://codeforces.com/contest/1698/problem/B
// Memory Limit: 256 MB
// Time Limit: 1000 ms
//
// Powered by CP Editor (https://cpeditor.org)

#include <iostream>
#include <vector>
#include <map>
#include <cstring>
#include <queue>
#include <math.h>
#include <set>
#include <stack>
#include <algorithm>
using namespace std;
#define IOS  ios::sync_with_stdio(false);
#define CIT  cin.tie(0);
#define COT  cout.tie(0);

#define ll long long
#define x first
#define y second
#define pb push_back
#define endl '\n'
#define all(x) (x).begin(),x.end()
#define Fup(i,a,b) for(int i=a;i<=b;i++)
#define Fde(i,a,b) for(int i=a;i>=b;i--)
#define NO cout<<"NO"<<endl;
#define YES cout<<"YES"<<endl;
#define cer(a) cerr<<#a<<'='<<(a)<<" @ line "<<__LINE__<<" "<<endl
typedef priority_queue<int,vector<int>,greater<int>>  Pri_m;
typedef pair<int,int> pii;
typedef vector<int> VI;
map<int,int> mp;
const int N = 2e5+10,INF = 0x3f3f3f3f;
const double eps = 1e-5;


struct node{
    int to,val;
};


int col[N],n;
int e[N],ne[N],h[N],idx;
int sz[N],son[N];


void add(int a,int b){
    e[idx] = b ,ne[idx] = h[a] , h[a] = idx++;
}
//树链剖分
void dfs(int u,int fa){
    sz[u] = 1;
    for(int i =  h[u];~i;i=ne[i]){
        int j = e[i];
        if(j == fa) continue;
        dfs(j,u);
        sz[u] += sz[j];
        if(sz[j] > sz[son[u]]) son[u] = j;
    }
}
int cnt[N];//某颜色在当前子树中的数量 
ll ans[N],sum;//答案数组 , 累加计算当前子树的答案
int  flag,maxn;//标记重儿子,更新最大值
//统计
void count(int u,int fa,int val){
    cnt[col[u]] += val;

    if(cnt[col[u]] > maxn){
        maxn = cnt[col[u]];
        sum = col[u];
    }else if(cnt[col[u]] == maxn) sum += col[u];

    for(int i = h[u];~i;i=ne[i]){
        int j = e[i];
        if(j == fa || j == flag) continue;
        count(j,u,val);

    }
}
void dfs(int u,int fa,bool keep){
    //轻儿子 算答案删除贡献
    for(int i  =h[u] ;~i ; i=ne[i]){
        int j = e[i];

        if(j == fa || j == son[u]) continue;

        dfs(j,u,false);//false 删除贡献

    }

    // 搞重儿子
    if(son[u]){
        dfs(son[u],u,true);
        flag = son[u];
    }

    // 暴力统计所有轻儿子的贡献
    count(u,fa,1);
    flag = 0 ;
    ans[u] = sum;

    if(!keep){
        count(u,fa,-1);
        sum = maxn = 0 ;

    }
}
void solve(){
    memset(h,-1,sizeof h);
    cin>>n;
    Fup(i,1,n) cin>>col[i];
    Fup(i,1,n-1){
        int u,v;cin>>u>>v;
        add(u,v);
        add(v,u);
    }
    dfs(1,0); //树剖找重儿子
    dfs(1,0,0);

    Fup(i,1,n) cout<<ans[i] <<" ";

}

int main(){
    //int t;cin>>t;while(t--)
    solve();
    return 0 ;
}