Consider managing the set of colors of the balls in each box, and updating them according to the operations.

If you simulate the operations just as given in the input, in each query the size of set of colors of the balls in box \(b\) increases at most by that of box \(a\) for each query, which may lead to TLE (Time Limit Exceeded).

Instead, we make use of the property that the operation in a query is equivalent to move all balls in box \(b\) into box \(a\), and swapping the two boxes.

Consider the following strategy: we compare the number of balls in boxes \(a\) and \(b\), and move all the balls in the one with fewer into the more, and swap them if needed.

This way, every time a ball is moved, the box that it belongs contains at least twice as many balls as before, so each ball is moved at most \(O(\log N)\) times, for a total of \(O(N \log N)\) moves.

One can swap boxes in constant time by managing indices appropriately, or using std::swap or std::move.

In the following implementation, boxes are compared by the number of distinct colors of balls instead of the number of balls themself, but the complexity can be similarly evaluated.

Sample code

#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]);
		}
	}
}