Codeforces Round #221 (Div. 1) D. Tree and Queries

問題

codeforces.com

解法

前使ったばかりの, データ構造をマージする一般的なテクを使います。

M[v][color] = (頂点 v 以下の subtree において, 色が color の頂点がいくつあるか)
N[v][num] = (頂点 v 以下の subtree において, 頂点数が num 以上ある color が何種類あるか)
という２つのデータ構造をマージしていきます。

size[v] = (v 以下の部分木にある頂点の数) を基にマージすると, 各頂点に対して, マージして移動する回数が O(log N) に抑えられるので, 全体計算量は O(N log N) になりますね。

const int MAXN = 100010;
int C[MAXN], ans[MAXN], size[MAXN];
vector<int> N[MAXN], G[MAXN];
map<int, int> M[MAXN];
vector<pii> Q[MAXN];

void merge(int v, int u) {
    if (size[v] < size[u]) {
        swap(size[v], size[u]);
        swap(M[v], M[u]);
        swap(N[v], N[u]);
    }
    size[v] += size[u];
    for (auto p : M[u]) {
        int x = M[v][p.first];
        for (int i = x+1; i <= x+p.second; i++) {
            if (N[v].size() < i) N[v].push_back(1);
            else N[v][i-1]++;
        }
        M[v][p.first] += p.second;
    }
}

void dfs(int v, int p) {
    M[v][C[v]]++;
    N[v].push_back(1);
    size[v] = 1;
    for (int ch : G[v]) {
        if (ch==p) continue;
        dfs(ch, v);
        merge(v, ch);
    }
    for (pii p : Q[v]) {
        if (p.first <= N[v].size()) ans[p.second] = N[v][p.first-1];
    }
}

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    int n, m;
    cin >> n >> m;
    for (int i = 0; i < n; i++) {
        cin >> C[i];
    }
    for (int i = 0; i < n-1; i++) {
        int u, v;
        cin >> u >> v;
        u--; v--;
        G[u].push_back(v);
        G[v].push_back(u);
    }
    for (int i = 0; i < m; i++) {
        int v, k;
        cin >> v >> k;
        Q[v-1].emplace_back(k, i);
    }
    dfs(0, -1);
    for (int i = 0; i < m; i++) cout << ans[i] << endl;
    return 0;
}