krsu 3403 (not my code)

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

#define ar array
typedef long long ll;

const int N = 5e5 + 5;
const int M = 20;
const int B = 700;
vector<int> edges[N], p, qq[N], sz[N];
int t[N], res[N], d[N], tin[N], tout[N];
int mn[N + N][M], b[N], v[N], dis[N];

void dfs(int u, int par = -1) {
    tin[u] = p.size();
    p.push_back(d[u]);
    tout[u] = tin[u];
    for (auto x : edges[u]) {
        if (x == par)
            continue;
        d[x] = d[u] + 1;
        dfs(x, u);
        tout[u] = p.size();
        p.push_back(d[u]);
    }
}

signed main() {
    ios::sync_with_stdio(0);
    cin.tie(0);

    int n, m;
    cin >> n >> m;
    for (int i = 1; i <= n; i++) {
        cin >> t[i];
        sz[t[i]].push_back(i);
    }
    for (int i = 1; i < n; i++) {
        int a, b;
        cin >> a >> b;
        edges[a].push_back(b);
        edges[b].push_back(a);
    }
    dfs(1);
    for (int i = 0; i < (int)p.size(); i++) {
        mn[i][0] = p[i];
    }

    for (int j = 1; j < M; j++) {
        for (int i = 0; i + (1 << (j - 1)) < (int)p.size(); i++) {
            mn[i][j] = min(mn[i][j - 1], mn[i + (1 << (j - 1))][j - 1]);
        }
    }

    auto get = [&](int a, int b) {
        if (tin[a] > tin[b])
            swap(a, b);
        if (tin[a] <= tin[b] && tin[b] <= tout[a]) {
            return d[b] - d[a];
        }
        int l = tout[a], r = tin[b];
        int lg = __lg(r - l + 1);
        return d[a] + d[b] - 2 * min(mn[l][lg], mn[r - (1 << lg) + 1][lg]);
    };

    int q;
    cin >> q;
    for (int i = 0; i < q; i++) {
        cin >> v[i] >> b[i];
        qq[b[i]].push_back(i);
    }

    for (int i = 1; i <= m; i++) {
        if ((int)sz[i].size() > B) {
            queue<int> q;
            memset(dis, 63, sizeof dis);
            for (auto x : sz[i])
                q.push(x), dis[x] = 0;
            while (!q.empty()) {
                int u = q.front();
                q.pop();
                for (auto x : edges[u]) {
                    if (dis[x] > dis[u] + 1) {
                        dis[x] = dis[u] + 1;
                        q.push(x);
                    }
                }
            }

            for (auto x : qq[i]) {
                res[x] = dis[v[x]];
            }
        } else {
            for (auto j : qq[i]) {
                res[j] = n;
                for (auto x : sz[i]) {
                    res[j] = min(res[j], get(x, v[j]));
                }
            }
        }
    }

    for (int i = 0; i < q; i++)
        cout << res[i] << "\n";
}

/*

6 5
3 5 1 1 4 2
3 6
2 3
4 3
6 1
3 5
6
3 5
2 3
2 1
4 4
1 1
4 2

*/