概要

この問題は、生徒のグループを効率的に管理し、グループの合併とグループのサイズのクエリに高速に答える問題です。

考察

素朴なアプローチでは、グループの合併を行うたびにすべての生徒の所属を更新すると、1回の合併に \(O(N)\) の時間がかかり、全体で \(O(NQ)\) となってしまい、制約条件 \(N, Q \leq 2 \times 10^5\) では実行時間が間に合いません。そこで、効率的にグループを管理するデータ構造が必要になります。

アルゴリズム

この問題はUnion-Find（Disjoint Set Union, DSU） データ構造を用いることで効率的に解くことができます。Union-Findは以下の2つの操作を高速に行えます： - find(x): 要素 \(x\) が属するグループの代表（ルート）を求める - union(a, b): 要素 \(a\) と \(b\) が属するグループを合併する

さらに、各グループのサイズ（要素数）を管理する配列 size を追加で保持します。これにより、グループのサイズのクエリにはルートのサイズを返すだけで答えることができます。

実装では、以下の2つの高速化手法を用いています： 1. 経路圧縮（Path Compression）: find 操作時に、調べたノードを直接ルートに接続し、木の高さを抑える 2. union by size: 合併時に、サイズが小さい木を大きい木に結合し、木の高さを抑える

これらの工夫により、各操作はほぼ定数時間 \(O(\alpha(N))\)（\(\alpha\) はアッカーマン関数の逆関数）で実行できます。

計算量

• 時間計算量: \(O(Q \cdot \alpha(N))\)
• 空間計算量: \(O(N)\)

実装のポイント

• クラスを使ってUnion-Findを実装すると、コードが整理されやすい

• 経路圧縮とunion by sizeを組み合わせることで効率化

• グループのサイズを管理する配列 size を用意し、ルートノードのサイズを保持

• 入力処理を高速化するため、sys.stdin.read で一括読み込みを行う

• 出力もリストに蓄えて最後にまとめて出力することで高速化

  ソースコード

class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n+1))
        self.size = [1] * (n+1)
    
    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]
    
    def union(self, a, b):
        root_a = self.find(a)
        root_b = self.find(b)
        if root_a == root_b:
            return
        if self.size[root_a] < self.size[root_b]:
            root_a, root_b = root_b, root_a
        self.parent[root_b] = root_a
        self.size[root_a] += self.size[root_b]

def main():
    import sys
    input = sys.stdin.read
    data = input().split()
    
    n = int(data[0])
    q = int(data[1])
    uf = UnionFind(n)
    
    index = 2
    output_lines = []
    for _ in range(q):
        t = data[index]; index += 1
        if t == '1':
            a = int(data[index]); b = int(data[index+1]); index += 2
            uf.union(a, b)
        else:
            x = int(data[index]); index += 1
            root = uf.find(x)
            output_lines.append(str(uf.size[root]))
    
    print("\n".join(output_lines))

if __name__ == "__main__":
    main()

この解説は deepseekv3 によって生成されました。