公式

F - 友達の輪 / Circle of Friends 解説 by admin

Claude 4.6 Opus (Thinking)

概要

グループの合併(union)とグループの人数の問い合わせ(find + size)を効率的に処理する問題です。Union-Find(素集合データ構造)を使うことで高速に解くことができます。

考察

素朴なアプローチとその問題点

各生徒がどのグループに属するかを配列で管理し、合併のたびに片方のグループの全員のグループ番号を書き換える方法が考えられます。しかし、この方法では1回の合併に最悪 \(O(N)\) かかり、\(Q\) 回の指示に対して全体で \(O(NQ)\) となります。\(N, Q\) が最大 \(2 \times 10^5\) のとき、これは \(4 \times 10^{10}\) 回の操作となり、制限時間に間に合いません。

重要な気づき

この問題は典型的な「要素のグループ分けを動的に管理する」問題であり、Union-Find(素集合データ構造) がぴったり当てはまります。Union-Find を使えば、合併と検索をほぼ \(O(1)\) で行えます。

アルゴリズム

Union-Find は、各グループを木構造で表現するデータ構造です。

  • parent[x]: 生徒 \(x\) の親を記録。根(ルート)がグループの代表者。
  • rank[x]: 木の高さの上限を記録(union by rank に使用)。
  • size[x]: \(x\) が根のとき、そのグループの人数を記録。

主要操作

  1. find(x): 生徒 \(x\) が属するグループの根を返す。途中で 経路圧縮(path splitting)を行い、次回以降の検索を高速化する。
  2. union(x, y): 生徒 \(x\)\(y\) のグループを合併する。union by rank(ランクが小さい木を大きい木の下に付ける)を行い、木の高さの増大を抑える。合併時に size を更新する。
  3. get_size(x): find(x) で根を求め、その根の size を返す。

具体例

\(N = 5\) で、以下の指示が来た場合を考えます。

  1. 1 1 2 → 生徒1と2を合併 → グループ: {1,2}, {3}, {4}, {5}
  2. 1 3 4 → 生徒3と4を合併 → グループ: {1,2}, {3,4}, {5}
  3. 2 1 → 生徒1のグループの人数 → 2 を出力
  4. 1 2 3 → 生徒2と3のグループを合併 → グループ: {1,2,3,4}, {5}
  5. 2 4 → 生徒4のグループの人数 → 4 を出力

計算量

  • 時間計算量: \(O((N + Q) \cdot \alpha(N))\)
    • \(\alpha(N)\) は逆アッカーマン関数で、実用上はほぼ定数(\(5\) 以下)です。
    • 経路圧縮と union by rank を併用することで、\(Q\) 回の操作全体でほぼ \(O(N + Q)\) となります。
  • 空間計算量: \(O(N)\)
    • parent, rank, size の3つの配列をそれぞれ \(N+1\) 要素分保持します。

実装のポイント

  • 1-indexed への対応: 生徒の番号が \(1\) から \(N\) なので、UnionFind(N + 1) としてサイズ \(N+1\) で初期化し、インデックス \(0\) は未使用としています。

  • 経路圧縮の実装: self.parent[x] = self.parent[self.parent[x]] により、find のたびに木を平坦化しています(path splitting と呼ばれる手法)。

  • 出力の一括処理: 指示2のたびに print するのではなく、結果をリストに溜めて最後に '\n'.join(out) で一括出力することで、I/O のオーバーヘッドを削減しています。

  • 高速な入力: sys.stdin.readlineinput に代入して、Python の標準入力を高速化しています。

    ソースコード

import sys
input = sys.stdin.readline

class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))
        self.rank = [0] * n
        self.size = [1] * n

    def find(self, x):
        while self.parent[x] != x:
            self.parent[x] = self.parent[self.parent[x]]
            x = self.parent[x]
        return x

    def union(self, x, y):
        rx, ry = self.find(x), self.find(y)
        if rx == ry:
            return
        if self.rank[rx] < self.rank[ry]:
            rx, ry = ry, rx
        self.parent[ry] = rx
        self.size[rx] += self.size[ry]
        if self.rank[rx] == self.rank[ry]:
            self.rank[rx] += 1

    def get_size(self, x):
        return self.size[self.find(x)]

def main():
    N, Q = map(int, input().split())
    uf = UnionFind(N + 1)
    out = []
    for _ in range(Q):
        query = list(map(int, input().split()))
        if query[0] == 1:
            a, b = query[1], query[2]
            uf.union(a, b)
        else:
            x = query[1]
            out.append(str(uf.get_size(x)))
    print('\n'.join(out))

main()

この解説は claude4.6opus-thinking によって生成されました。

投稿日時:
最終更新: