distinct colors cses

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

const ll nmax = 1e9+7;
const ll nmax2 = 998244353;

int main(){
    std::ios_base::sync_with_stdio(false);
    std::cin.tie(0);
    std::cout.tie(0);

    ll n;
    std::cin >> n;

    ll c[n + 1];
    for (ll i = 1; i <= n; i++) std::cin >> c[i];

    std::vector<ll> adj[n + 1];
    for (ll i = 1, x, y; i < n; i++){
        std::cin >> x >> y;
        adj[x].push_back(y);
        adj[y].push_back(x);
    }

    std::set<ll> s[n + 1];
    std::function<void(ll, ll)> dfs = [&] (ll u, ll p){
        std::vector<std::pair<ll, ll>> child;
        for (auto v: adj[u]){
            if (v != p){
                dfs(v, u);
                child.push_back({s[v].size(), v});
            }
        }
        if (child.size()){
            std::sort(child.begin(), child.end());
            std::reverse(child.begin(), child.end());
        }
        if (child.size()){
            s[u] = s[child[0].second];
            if (child.size() > 1){
                for (ll i = 1; i < child.size(); i++){
                    for (auto v: s[child[i].second]){
                        s[u].insert(v);
                    }
                }
            }
        }
        s[u].insert(c[u]);
    };

    dfs(1, 0);

    for (ll i = 1; i <= n; i++){
        std::cout << s[i].size() << ' ';
    }
}