Code : Distinct Colors

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

class Tree {
  public:
    int n;
    vector<vector<int>> adj;
    vector<set<int>> uniq;
    vector<int> a;

    Tree(int n) {
        this->n = n;
        adj.resize(n);
        uniq.resize(n);
        a.resize(n);
    }

    void dfs(int src, int par) {
        uniq[src].insert(a[src]);
        for (auto child : adj[src]) {
            if (child == par) {
                continue;
            }
            dfs(child, src);
            for (auto &ele : uniq[child]) {
                uniq[src].insert(ele);
            }
        }
    }
};

void solve() {
    int n;
    cin >> n;
    Tree t(n);

    for (int i = 0; i < n; i++) {
        cin >> t.a[i];
    }
    for (int i = 0; i < n - 1; i++) {
        int x, y;
        cin >> x >> y;
        x--;
        y--;
        t.adj[x].push_back(y);
        t.adj[y].push_back(x);
    }

    t.dfs(0, -1);
    for (int i = 0; i < n; i++) {
        cout << t.uniq[i].size() << " ";
    }
    cout << "\n";
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    solve();
    return 0;
}

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

class Tree {
  public:
    int n;
    vector<vector<int>> adj;
    vector<set<int>> uniq;
    vector<int> a, uniq_sz;

    Tree(int n) {
        this->n = n;
        adj.resize(n);
        uniq.resize(n);
        a.resize(n);
        uniq_sz.resize(n);
    }

    void dfs(int src, int par) {
        uniq[src].insert(a[src]);
        for (auto child : adj[src]) {
            if (child == par) {
                continue;
            }
            dfs(child, src);

            if (uniq[child].size() > uniq[src].size()) {
                swap(uniq[child], uniq[src]);
            }
            for (auto &ele : uniq[child]) {
                uniq[src].insert(ele);
            }
        }
        uniq_sz[src] = uniq[src].size();
    }
};

void solve() {
    int n;
    cin >> n;
    Tree t(n);

    for (int i = 0; i < n; i++) {
        cin >> t.a[i];
    }
    for (int i = 0; i < n - 1; i++) {
        int x, y;
        cin >> x >> y;
        x--;
        y--;
        t.adj[x].push_back(y);
        t.adj[y].push_back(x);
    }

    t.dfs(0, -1);
    for (int i = 0; i < n; i++) {
        cout << t.uniq_sz[i] << " ";
    }
    cout << "\n";
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    solve();
    return 0;
}