箱に入っているボールの色の集合を管理しておき、操作に応じて各集合を更新していくことを考えます。

問題文で与えられた通りに操作を行うと、各クエリで箱 \(b\) に入っているボールの色の集合の要素数は最大で箱 \(a\) に入っているボールの数だけ増える可能性があり、このまま実装すると TLE してしまいます。

そこで、各クエリでする操作は箱 \(b\) のボールをすべて箱 \(a\) に移し、箱 \(a\) と箱 \(b\) を交換することと等しいことを利用します。

箱 \(a\) と箱 \(b\) のうち箱に入っているボールの色の個数が少ない方を多い方に移し、必要に応じて箱を交換することを考えます。

そのようにすると、各ボールについて、そのボールが移されたときに同じ箱に入っているボールの個数は \(2\) 倍以上になるため \(O(\log N)\) 回しかボールの移動は発生せず、全体としてボールの移動は \(O(N \log N)\) 回しか発生しないことがわかります。

箱の交換は index を適切に管理する、あるいは C++ における std::swap や std::move などを利用することで定数時間で行えます。

以下の実装例では箱に入っているボールの個数ではなく種類数でどちらを移すか比較していますが、同様の原理で計算量が評価できます。

実装例

#include <iostream>
#include <vector>
#include <set>
using namespace std;

int main() {
	int n, q;
	cin >> n >> q;
	vector<set<int>> st(n);
	for (int i = 0; i < n; i++) {
		int c;
		cin >> c;
		st[i].insert(c);
	}
	while (q--) {
		int a, b;
		cin >> a >> b;
		a--; b--;
		if (st[a].size() < st[b].size()) {
			for (int i : st[a]) st[b].insert(i);
			st[a].clear();
			cout << st[b].size() << '\n';
		}
		else {
			for (int i : st[b]) st[a].insert(i);
			st[b].clear();
			cout << st[a].size() << '\n';
			swap(st[a], st[b]);
		}
	}
}