各集団について「各クラスに属する生徒が何人いるか」をmapやdictのようなデータ構造で持つことを考えます。集団そのものの管理にはUnion-Findを用います。

クエリ2に対しては、生徒 \(x\) がどの集団に属しているかをUnion-Findで求め、該当する集団のmap/dictを調べることで \(O(\log N)\) で答えることができます。

クエリ1に対しては、「人数が少ない方の集団を、人数が多い方の集団に吸収させる」として、2つのmap/dictをマージします。map/dictの各要素を、一方から他方へ愚直に移します。

実装例(C++)

//map<int,int>であるmap1 の中身を map2に移す
for(auto [key, value]:map1) map2[key]+=value;

実装例(Python)

#defaultdictであるdict1の中身をdict2に移す
for key, value in dict1.items():
    dict2[key]+=value;

この操作には、map1/dict1の要素数を \(k\) として \(O(k\log N)\) の時間がかかるため、最悪計算量は \(O(N\log N)\) です。 しかし、各生徒が属している集団の大きさは、移動が行われるごとに2倍以上になるため、各生徒について移動が行われる回数は高々 \(O(\log N)\) 回です。したがってクエリ1全てについての計算量の合計は高々 \(O(N(\log N)^2)\) になります。

以上により、この問題を \(O(N(\log N)^2+Q\log N)\) で解くことができました。

解答例(C++)

解答例(Python)

一般に、今回のように「小さい方から大きい方へ移す」とすることで最悪計算量が改善されることは多くあります。（Union-Find の Union by size による実装など）