Codeforces Round #221 Div.1 D Tree and Queries

問題文

http://codeforces.com/contest/375/problem/D

問題概要

頂点に色が付いたサイズ N の根付き木がある. 根は 1. 以下の Q 個のクエリに答えよ.

• v k : 頂点 v の部分木に k 個以上含まれるような色の種類数を出力

解法

前に解いたときは木上 Mo で unordered_map を使って解いたけど, 若干不正AC感があったので(0.92s / 1s)

これは実は DSU on Tree というのを使うともっと速く解けます.
DSU on Tree について詳しくはこれ
codeforces.com

イメージは dfs で愚直にやるけど, centroid-path でつながった子のときだけ計算結果を持ち越してくる感じ.
クエリは前処理で各頂点に持たせておいて最後にまとめて答えます.

で, どういう状態を持たせるのかが重要で, 適当に各色のカウントを持つだけでは各クエリの計算に O(N) かかるのでダメ.
各色のカウントをさらに unordered_map とかで持てば, 一応追加/削除 O(1) で各クエリに O(√N) *1で答えられるけどまだ遅い.(ちなみに前に解いたときはこれを使った)
一番良いのは各色のカウントと, 各 k について k 個以上含まれる色の種類数を持たせます.(k=0は適当でいい)
更新ができるのか一瞬不安になりますが, 実は簡単で追加/削除 O(1), クエリ O(1) できます.

これで全体で O(NlogN + Q) になるので安心.

ソースコード

http://codeforces.com/contest/375/submission/40633975

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

const int MAX = 1e5;

vector<int> G[MAX];

int size[MAX];

void build(int v, int prev) {
    size[v] = 1;
    for (auto& p : G[v]) if (p == prev) {
        if (G[v].back() != p) swap(G[v].back(), p);
        G[v].pop_back();
        break;
    }
    for (auto to : G[v]) {
        build(to, v);
        size[v] += size[to];
    }
    for (auto& to : G[v]) {
        if (size[to] > size[G[v].back()]) {
            swap(to, G[v].back());
        }
    }
}

struct query {
    int id, k;
    query(int i, int k_) : id(i), k(k_) {}
};

vector<query> qs[MAX];
int res[MAX];

vector<int> vs[MAX];

int c[MAX];
int cnt[MAX];
int dat[MAX + 1];

void add(int v) {
    ++cnt[c[v]];
    ++dat[cnt[c[v]]];
}

void del(int v) {
    --dat[cnt[c[v]]];
    --cnt[c[v]];
}

void dfs(int v, bool keep = false) {
    int hv = G[v].empty() ? -1 : G[v].back();
    for (auto to : G[v]) {
        if (to != hv) {
            dfs(to);
        }
        else {
            dfs(hv, true);
            vs[v] = move(vs[hv]);
        }
    }
    add(v);
    vs[v].push_back(v);
    for (auto to : G[v]) if (to != hv) {
        for (auto ver : vs[to]) add(ver);
        copy(vs[to].begin(), vs[to].end(), back_inserter(vs[v]));
    }
    for (auto& q : qs[v]) {
        res[q.id] = dat[q.k];
    }
    if (!keep) {
        for (auto ver : vs[v]) del(ver);
    }
}

int main()
{
    ios::sync_with_stdio(false), cin.tie(0);
    int n, m;
    cin >> n >> m;
    for (int i = 0; i < n; i++) {
        cin >> c[i]; --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; --v;
        qs[v].emplace_back(i, k);
    }
    build(0, -1);
    dfs(0);
    for (int i = 0; i < m; i++) {
        printf("%d\n", res[i]);
    }
    return 0;
}

蛇足

DSU on Tree を覚えたので使ってみました.
おそらくこういう問題で DSU on Tree を使う時はマージテクでも同じようにできそう.
個人的にはマージテクより実装は少し楽に感じました.
あと, 部分木に対する木上 Mo とできることがほぼ変わらなさそうなので, こういう問題でなら完全上位互換っぽい.

*1:全ての色のカウントの総和が高々 N だから, カウント数の種類は O(√N) になるため