零食采购 参考代码

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

typedef long long LL;
const int maxn = 100000 + 100;
const int Log = 20;
int n, q, u, v, s, t, cnt;
int c[maxn], id[maxn];
int st[maxn][Log], bitMask[maxn][Log];
vector<int> G[maxn];

void dfs(int root, int fa) {
    id[root] = ++cnt;
    st[root][0] = fa;
    bitMask[root][0] = 1 << c[root];
    for (int i = 1; i < Log; ++i) {
        st[root][i] = st[st[root][i - 1]][i - 1];
        bitMask[root][i] = bitMask[root][i - 1] | bitMask[st[root][i - 1]][i - 1];
    }
    for (int pos : G[root]) {
        if (pos == fa) {
            continue;
        }
        dfs(pos, root);
    }
}

int solve(int s, int t) {
    if (s == t) {
        return 1;
    }
    if (id[s] > id[t]) {
        swap(s, t);
    }
    int mask = 0;
    for (int i = Log - 1; i >= 0; --i) {
        if (id[st[t][i]] > id[s]) {
            mask |= bitMask[t][i];
            t = st[t][i];
        }
    }
    mask |= bitMask[t][0];
    t = st[t][0];
    for (int i = Log - 1; i >= 0; --i) {
        if (id[st[s][i]] > id[t]) {
            mask |= bitMask[s][i];
            s = st[s][i];
        }
    }
    mask |= bitMask[s][0];
    return __builtin_popcount(mask | (1 << c[t]));
}

int main() {
#ifdef ExRoc
    freopen("test.txt", "r", stdin);
#endif
    ios::sync_with_stdio(false);

    cin >> n >> q;
    for (int i = 1; i <= n; ++i) {
        cin >> c[i];
        --c[i];
    }
    for (int i = 1; i < n; ++i) {
        cin >> u >> v;
        G[u].push_back(v);
        G[v].push_back(u);
    }
    dfs(1, 1);
    while (q--) {
        cin >> s >> t;
        cout << solve(s, t) << endl;
    }

    return 0;
}