We will consider storing the number of students belonging each class for each group with data structures like map or dict. We use Union-Find to manage the groups themselves.

For a Query 2, we can use Union-Find to find the group that Student \(x\) belongs, and inspect the map/dict for the corresponding group, so that we can answer the query in an \(O(\log N)\) time.

For a query 1, we merge the two maps/dicts so that “the group with the fewer people is merged into that with the more people.” We transfer each element of map/dict naively, one to the other.

Sample Code (C++)

//move the content of map1, which is a map<int,int>, to map2
for(auto [key, value]:map1) map2[key]+=value;

Sample Code(Python)

#move the content of dict1, which is a defatuldict, to map2
for key, value in dict1.items():
    dict2[key]+=value;

This operation costs an \(O(k\log N)\) time, where \(k\) is the number of elements in map1/dict1, so the worst computation time is \(O(N\log N)\). However, since every move at least doubles the size of group that each student belongs, the number of transfers applied to each student is at most \(O(\log N)\). Therefore, the total computation time for all Queries 1 is at most \(O(N(\log N)^2)\).

Therefore, the problem has been solved in a total of \(O(N(\log N)^2+Q\log N)\) time.

Sample Code (C++)

Sample Code (Python)

In general, it is often the case that the worst computation time can be improved by “transferring from the smaller to the larger.” (As in the implementation of “Union by size” in Union-Find)