Code : Confuzzle

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

const int inf = 1e9;

class Tree {
  public:
    int n;
    vector<vector<int>> adj;
    vector<int> a, depth;
    vector<map<int, int>> min_depth;
    int ans;

    // min_depth[i][c] is the minimum depth of the color c in the subtree of
    // the the i-th node.

    Tree(int n) {
        this->n = n;
        adj.resize(n);
        a.resize(n);
        depth.resize(n);
        min_depth.resize(n);
        ans = inf;
    }

    void dfs(int src, int par) {
        if (par != -1) {
            depth[src] = depth[par] + 1;
        }

        min_depth[src][a[src]] = depth[src];
        for (auto child : adj[src]) {
            if (child == par) {
                continue;
            }
            dfs(child, src);

            // Parent is always the bigger set.
            if (min_depth[child].size() > min_depth[src].size()) {
                swap(min_depth[child], min_depth[src]);
            }

            for (auto &[c, d] : min_depth[child]) {
                if (min_depth[src].find(c) == min_depth[src].end()) {
                    min_depth[src][c] = d;
                } else {
                    int have = min_depth[src][c];
                    ans = min(ans, have + d - 2 * depth[src]);
                    min_depth[src][c] = min(have, d);
                }
            }
        }
    }
};

void solve() {
    int n;
    cin >> n;
    Tree t(n);
    for (int i = 0; i < n; i++) {
        cin >> t.a[i];
        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);

    cout << t.ans << "\n";
}

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

    solve();
    return 0;
}