This is an application of the “union-by-size / small-to-large heuristic”. Simply put, if the balls in box \(a\) are never more than the balls in box \(b\), you could simply move the balls one by one, and each ball will not move more than \(O(\log N)\) times, as each time it moves, the number of balls in the new box is at least twice that of the old box, and there are no more than \(N\) balls in total.

To solve the original problem, if box \(a\) has more balls than box \(b\), we could swap the two boxes before merging. Most languages have set data structures that store the bulk of their data on the “memory heap”, while having only a small amount of data, like pointers and size, in the array that contains the sets; so swapping merely swaps the pointer/size data and not the individual elements of the sets; so we could accomplish this swap in \(O(1)\).

Note that if we store the colors instead of the balls directly (as we can then query the size of the set directly), some entries could be “lost” during merging; however, this would not slow down the solution, as it means we might move fewer entries than expected by the heuristic, but never more.

If we use ordered sets to store the colors, each ball move could cost \(O(\log N)\), so the problem is solved in \(O(N \log ^2 N)\). It can also be solved in \(O(N \log N)\) using an unordered set (hash set), though be sure to use a well-distributed / randomized hash function.

Link to similar problem: ABC 279 F: BOX