[Codeforces] E. Lomsat gelral

Author            : Dipu Kumar Mohanto 
                    CSE, Batch - 6
                    BRUR.
Problem Statement : E. Lomsat gelral
Source            : Codeforces
Category          : Graph Theory
Algorithm         : DSU on Tree
Verdict           : Accepted

#include "bits/stdc++.h"
#include "ext/pb_ds/assoc_container.hpp"
#include "ext/pb_ds/tree_policy.hpp"
 
using namespace std;
using namespace __gnu_pbds;
 
#define ll            long long
#define eb            emplace_back
 
static const int maxn = 2e5 + 5;
 
template <typename T> using orderSet = tree <T, null_type, greater <T>, rb_tree_tag, tree_order_statistics_node_update>;
 
orderSet <ll> X;
vector <int> graph[maxn];
ll color[maxn], sizeSubTree[maxn], cnt[maxn];
bool big[maxn];
ll ans[maxn];
ll occSum[maxn];
 
void dfs(int u = 1, int p = -1)
{
    sizeSubTree[u] = 1;
    for (int v : graph[u])
    {
        if (v == p) continue;
        dfs(v, u);
        sizeSubTree[u] += sizeSubTree[v];
    }
}
 
void Add(ll num)
{
    X.insert(num);
}
 
void Erase(ll num)
{
    X.erase(X.lower_bound(num));
}
 
void update(int u, int p, ll x)
{
    // first erase the previous color
    Erase(cnt[ color[u] ]);
    occSum[ cnt[ color[u] ] ] -= color[u];
    // increment color size
    cnt[ color[u] ] += x;
    // Add to ordered set
    Add(cnt[ color[u] ]);
    occSum[ cnt[ color[u] ] ] += color[u];
 
    for (int v : graph[u])
    {
        if (v == p || big[v]) continue;
        update(v, u, x);
    }
}
 
ll getResult(int u)
{
    int f = *X.find_by_order(0);
    return occSum[f];
}
 
void dsu(int u = 1, int p = -1, bool keep = 0)
{
    int mx = -1, bigChild = -1;
    for (int v : graph[u])
    {
        if (v == p) continue;
        if (sizeSubTree[v] > mx) mx = sizeSubTree[v], bigChild = v;
    }
    for (int v : graph[u])
    {
        if (v == p || v == bigChild) continue;
        dsu(v, u, 0); // Run a dfs from small child and clear them from cnt
    }
    if (bigChild != -1)
    {
        dsu(bigChild, u, 1);
        big[bigChild] = 1; // bigchild marked as big and not cleared from cnt
    }
    update(u, p, 1);
    ans[u] = getResult(u);
    if (bigChild != -1) big[bigChild] = 0;
    if (keep == 0) update(u, p, -1);
}
 
int main()
{
    //freopen("in.txt", "r", stdin);
 
    int tnode;
    scanf("%d", &tnode);
    for (int i = 1; i <= tnode; i++) scanf("%lld", color+i);
    for (int e = 1; e < tnode; e++)
    {
        int u, v;
        scanf("%d %d", &u, &v);
        graph[u].eb(v);
        graph[v].eb(u);
    }
    dfs();
    dsu();
    for (int i = 1; i <= tnode; i++) printf("%lld ", ans[i]);
}