人生で初めて「マージテク」を学ぶなら、この問題！！！

問題へのリンク

問題概要

箱 $1, 2, \dots, N$ がある。箱 $i$ には最初色 $C_{i}$ のボールが入っている。以下の $Q$ 回のクエリに答えよ。

整数 $a, b$ が与えられる
箱 $a$ のボールをすべて箱 $b$ に移動させる
その後、箱 $b$ に入っているボールの色の種類数を答える

制約

$N, Q \le 2 \times 10^{5}$
$1 \le C_{i} \le N$

考えたこと

まさに「マージテク」の問題。

この問題のクエリを愚直に実行すると、一見 $O(Q N)$ の計算量を要するように見える。つまり、次のような実装が考えられる。

各箱に入っているボールの色の種類を vector<unordered_set<int>> 型の変数（box とする）で管理する
box[a] の要素を 1 つ 1 つ、box[b] に挿入していく

    while (Q--) {
        cin >> a >> b;
        a--, b--;
        for (auto v : box[a]) box[b].insert(v);
        box[a].clear();
        cout << box[b].size() << '\n';
    }

確かに、このままだと $O(QN)$ の計算量を要する。

しかし、実は次の工夫をするだけで、 $O(Q + N \log N)$ の計算量に改善できるのだ。

box[a] の要素を 1 つ 1 つ box[b] に挿入していく前に、もし

box[a].size() > box[b].size()

であるならば、swap(box[a], box[b]) をしておく

こうすることで、より少ないボールの移動でクエリを実行できる。まさかこのような小さな工夫が、計算量のオーダーレベルの改善をもたらすのは驚きだ！！

計算量が改善される理由

工夫後の実装において、各ボール $j$ について、何回箱を移動することになるかを考えてみよう。

ボール $j$ が箱 $a$ から箱 $b$ に移動するとき、box[a].size() <= box[b].size() であるとする（上述の工夫を施した場合の仮定）。

この仮定のもとでは、ボール $j$ の属する箱に入っているボールの色の種類数は、移動前後で 2 倍以上になっている。そのため、ボール $j$ の属する箱が入れ替わる回数は高々

$O(\log N)$ (回)

で抑えられるのだ。どのボールについても、このことが言えるので、クエリ全てに答えるのに要する計算量は $O(Q + N \log N)$ と評価できる。

コード

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

int main() {
    int N, Q, a, b, C;
    cin >> N >> Q;
    vector<unordered_set<int>> box(N);
    for (int i = 0; i < N; i++) {
        cin >> C;
        box[i].insert(C);
    }
    while (Q--) {
        cin >> a >> b;
        a--, b--;
        if (box[a].size() > box[b].size()) swap(box[a], box[b]);
        for (auto v : box[a]) box[b].insert(v);
        box[a].clear();
        cout << box[b].size() << '\n';
    }
}

けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 329 F - Colored Ball (1D, 水色, 500 点)

問題概要

制約

考えたこと

計算量が改善される理由

コード