DSU on tree (http://codeforces.com/contest/600/problem/E), C++ (gcc)

DSU on tree (http://codeforces.com/contest/600/problem/E)

#include <bits/stdc++.h>

using namespace std;
#define pb push_back

typedef long long int ll;

const int N = 1e5 + 1;

vector<int> g[N];

ll big[N] = {0}, sub[N], col[N], cnt[N] = {0}, sum[N] = {0}, ans[N] = {0};

void dfs(int s, int p){
    sub[s] = 1;
    for(auto &i : g[s])
        if(i != p){
            dfs(i, s), sub[s] += sub[i];
            if(sub[i] > sub[g[s][0]]) swap(i, g[s][0]);
        }
}

multiset<int> st;

void add(int s, int p, int v){
    int cur = col[s];
    if(cnt[cur] != 0) st.erase(st.find(cnt[cur])), sum[cnt[cur]] -= cur;
    cnt[cur] += v;
    if(cnt[cur] != 0) st.insert(cnt[cur]), sum[cnt[cur]] += cur;
    
    for(auto i : g[s])
        if(i != p && !big[i]) add(i, s, v);
}

void dfs_dsu(int s, int p, bool keep){
    int x = g[s][0];
    for(auto i : g[s])
        if(i != x && i != p) dfs_dsu(i, s, 0);
    if(x != p) dfs_dsu(x, s, 1), big[x] = 1;
        
    add(s, p, 1);
    ans[s] = sum[*st.rbegin()];
    if(x != p) big[x] = 0;
    
    if(!keep) add(s, p, -1);
}

int main(){
    ios_base :: sync_with_stdio(0);
    cin.tie(0);
    
    int n;
    cin >> n;
    for(int j = 1; j <= n; j++) cin >> col[j];
    for(int j = 0; j < n-1; j++){
        int u, v;
        cin >> u >> v;
        g[u].pb(v);
        g[v].pb(u);
    }
    dfs(1, 0);
    dfs_dsu(1, 0, 0);
    for(int j = 1; j <= n; j++)
        cout << ans[j] << " ";
    cout << endl;
    return 0;
}


run \| edit \| history \| help	0