問題

$1, 2, \cdots , n$ の番号がついた $n$ 個の箱があり、はじめ $i$ 番目の箱には色 ${c}_{i}$ のボールが1個入っている。

整数組 $(a, b)$ によって与えられる、以下の内容のクエリが $q$ 個与えられるので、これらを順に処理せよ。

箱 $a$ に入っているボールをすべて箱 $b$ に移し、その後箱 $b$ に入っているボールの色の種類数を出力する。

入力

まず最初の1行目には、整数 $n, q$ が与えられる。

その次の1行には、整数 ${c}_{i}$ が $n$ 個順番に与えられる。

その後の $q$ 行には、それぞれ $i$ 番目のクエリとして用いられる整数 $a, b$ が与えられる。

条件

$1 \leq n, q \leq 2 \times {10}^{5}$
$1 \leq {c}_{i} \leq n$
$1 \leq a, b \leq n$
$a \ne b$

出力

$q$ 行にわたり、各 $i$ 番目のクエリの答えを出力すること。

解法

単純にクエリを処理するために、 $n$ 個の箱のそれぞれについて、現在その箱に入っている色の種類を記録するためのものを用意します。

箱の中に入っているボールの色が重複する場合もあるため、色の記録のために vector を用いると、種類数をすぐに出すことができません。

そこで、 set *1 を用いて箱に入っている色の種類を記録していくことを考えます。 set では、その中に含まれるユニークな要素のみを管理しているため、.size() を出力するのみで色の種類数を出力することができます。

次に、箱 $a$ に入っているボールをすべて箱 $b$ に移す、という作業をどのように管理するのかということについて考えていきます。

箱 $a, b$ に入っている色の種類を管理している set をそれぞれ box[a], box[b] とすると、「 box[a] の中の要素をすべて box[b] に投入し、box[a] の中の要素を空にする」という単純な方法で行えます。

しかし、移し替える色の種類数が毎回 $n$ 種類ある場合について考えると、この単純な方法では計算量が $O(nq)$ となるので、実行時間制限には間に合いません。（今回の問題では制限が 4s とやや長めになっているので、間に合うかなと思って実験してみましたが、ダメでした*2。）

ここで、「箱 $a$ に入っているボールをすべて箱 $b$ に移し、その後箱 $b$ に入っている色の種類数を出力する」という作業は、「箱 $a$ と箱 $b$ に含まれる色の種類数を出力する」という作業と同義です。

従って、「箱 $a$ と箱 $b$ のうち、要素数の少ない方から多い方へとボールを移す」ということにより、移し替える作業についての時間のロスを防いだ上で、色の種類数を出力できるということが言えます。

一方で、このようにすると箱のインデックスが変わってしまいます。そのため、こちら側で仮のインデックスを用意しておき、クエリで与えられる本物の箱のインデックスを管理していくことを考えます。

例えば、「箱1に1種類、箱2に2種類、箱4に4種類の色のボールがそれぞれ入っている」という状況を考えます。（すべての色は異なるものとします。）

この仮定で、まずクエリ 4 1 が与えられるとします。

この際に、実際には「箱4から箱1に4種類のボールを移す」という作業が行われ、最終的には「箱1に5種類のボールがある」という状況になります。

しかし、箱4に入っている色の種類数が多いので、コード上では「box[1] のすべての要素を box[4] へと移す」という作業を行うことになります。これを行うと「box[1] には0種類、box[4] には5種類のボールが存在する」ということになり、実際の箱の番号と辻褄が合わなくなってしまいます。

この自体を防ぐには、移し替えの作業終了後に「箱1は box[4] を指しており、箱4は box[1] を指している」という対応関係を記録していくことが必要となります。

そして次に、クエリ 1 2 が与えられるとします。

このクエリでは、実際には「箱1から箱2へと5種類のボールを移す」という作業が行われ、最終的には「箱2に7種類のボールがある」という状況になります。

コード上の状況を整理すると、箱1を指しているものは box[4] であり、箱2を指しているものは box[2] です。

そして、 box[4] には5種類のボールがあり、 box[2] には2種類のボールがあるため、コード上では「box[2] のすべての要素を box[4] へと移す」という作業を行うことになります。

このことから、コード上では「box[4] に7種類のボールが存在する」ということになるため、移し替えの作業終了後に「箱1は box[2] を、箱2は box[4] をそれぞれ指している」という対応関係を記録する必要があります。

以上のような場合を実際に図示してみると、コード上でのインデックスの対応関係は以下のように変化していると言えます。

以上から、実際の移し替えと逆の移し替えをコード上で行うときについては、「対応するインデックスを入れ替える」という作業を行うことで、実際の箱のインデックスとコード上でのインデックスを管理することができます。

この方法により、与えられたクエリに対して「ボールの種類数が少ない方から多い方へ移し替える」ということができるようになるため、計算量を減らしてクエリに対しての出力を行えるようになります。

ソースコード

以下のようにして、 main() 関数の中に、インデックス管理を行いクエリを処理する部分を実装しました。

int n, q, c[int(2e5 + 5)];
int box_index[int(2e5 + 5)]; // コード上での箱のインデックスを表す
set<int> box[int(2e5 + 5)]; // 実際に入っている色を管理する set

int main() {
  cin >> n >> q;
  for (int i = 1; i <= n; i++) {
    cin >> c[i];
    box[i].insert(c[i]); // box[i] に色 c[i] が入っていることを set に記録する
    box_index[i] = i; // まずは初期状態として、箱 i のコード上のインデックスも i にする
  }

  // q 個のクエリの処理
  for (int i = 0; i < q; i++) {
    int a, b;
    cin >> a >> b;

    // 箱 a, b がコード上でどのインデックスに当たるのかを持つ
    int a_index = box_index[a], b_index = box_index[b];

    // もし箱 a の種類数が箱 b より多いならば、コード内でのインデックスを逆転させる
    if (box[a_index].size() > box[b_index].size()) {
      swap(a_index, b_index);
      swap(box_index[a], box_index[b]);
    }

    // box[a_index] 内のすべての要素を box[b_index] に投入する
    auto iter = box[a_index].begin();
    while (iter != box[a_index].end()) {
      box[b_index].insert(*iter);
      iter++;
    }

    // box[a_index] を空にする
    box[a_index] = {};

    // box[b_index] の要素数を出力する 
    cout << box[b_index].size() << endl;
  }

  return 0;
}