【題解】CSES 1139 Distinct Colors

#include <iostream>
#include <map>
#include <vector>
#include <algorithm>
#include <cmath>
using namespace std;
 
int n, k, a[200005], b[200005], cnt[200005], in[200005], out[200005], res[200005], ans, idx;
vector <int> v, e[200005];
vector <pair<int, pair<int, int> > > q;
 
bool cmp(pair<int, pair<int, int> > x, pair<int, pair<int, int> > y){
    if (x.second.first/k != y.second.first/k) return x.second.first/k < y.second.first/k;
    else return x.second.second < y.second.second;
}
void dfs(int x, int p){
    in[x] = idx;
    b[idx] = a[x];
    for (auto i:e[x]){
        if (i == p) continue;
        idx++;
        dfs(i, x);
    }
    out[x] = idx;
}
int id(int x){
    return lower_bound(v.begin(), v.end(), x)-v.begin();
}
void add(int x){
    if (cnt[x] == 0) ans++;
    cnt[x]++;
}
void remove(int x){
    if (cnt[x] == 1) ans--;
    cnt[x]--;
}
 
int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cin >> n;
    k = sqrt(n);
    for (int i = 1; i <= n; i++){
        cin >> a[i];
        v.push_back(a[i]);
    }
    sort(v.begin(), v.end());
    v.erase(unique(v.begin(), v.end()), v.end());
    for (int i = 1; i <= n; i++){
        a[i] = id(a[i]);
    }
    int x, y;
    for (int i = 1; i < n; i++){
        cin >> x >> y;
        e[x].push_back(y);
        e[y].push_back(x);
    }
    dfs(1, 0);
    for (int i = 1; i <= n; i++){
        q.push_back({i, {in[i], out[i]}});
    }
    sort(q.begin(), q.end(), cmp);
    int l = 0, r = 0;
    add(b[0]);
    for (auto i:q){
        while (i.second.first < l){
            l--;
            add(b[l]);
        }
        while (r < i.second.second){
            r++;
            add(b[r]);
        }
        while (l < i.second.first){
            remove(b[l]);
            l++;
        }
        while (i.second.second < r){
            remove(b[r]);
            r--;
        }
        res[i.first] = ans;
    }
    for (int i = 1; i <= n; i++){
        cout << res[i] << " ";
    }
}