公式

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

Qwen3-Coder-480B

概要

生徒のグループ分けを管理し、グループの結合とサイズ取得を効率的に行う問題です。

考察

この問題では、生徒を頂点、グループを連結成分とするグラフを考えることができます。
指示1(グループの合併)は、グラフ上で2つの連結成分を1つにすることに対応し、
指示2(グループのサイズ取得)は、ある頂点が含まれる連結成分の頂点数を求めることに対応します。

素朴な方法として、グループをリストや集合で管理し、合併のたびに全員を移動させるやり方がありますが、
最悪ケースでは1回の操作で \(O(N)\) かかり、全体で \(O(QN)\) となって間に合いません(TLE)。

そこで、高速にグループの合併とサイズ取得ができるデータ構造「Union-Find木(またはDisjoint Set Union)」を使用します。
Union-Find木では、各要素が親を持つ木構造を持ち、根を代表としてグループを管理します。
さらに、経路圧縮とrank/sizeによるマージ最適化を行うことで、ほぼ定数時間で操作が可能です。

また、指示2の出力を逐次printすると遅くなる可能性があるため、一旦結果をリストに保存して最後にまとめて出力する工夫がされています。

アルゴリズム

Union-Find木(Disjoint Set Union) を使用します。

  • 各生徒に対応するノードを持ち、初期状態ではすべてのノードが自身を親とする(つまり、各グループは1人だけ)。
  • unite(a, b):生徒 \(a\) と \(b\) が属するグループを統合。
    • 各ノードの根(代表)をfindで求め、異なる場合にのみ統合。
    • サイズが小さい方の木を大きい方に繋ぐ(union by size)ことで木の高さを抑える。
  • get_size(x):生徒 \(x\) が属するグループのサイズを返す。
    • 根までの経路を圧縮しつつ(find)、根のサイズを返す。

このアルゴリズムにより、グループの統合とサイズ取得を非常に高速に行えます。

計算量

  • 時間計算量: \(O(N + Q \cdot \alpha(N))\)
    (\(\alpha(N)\) はアッカーマン関数の逆関数で、実質定数とみなせる)
  • 空間計算量: \(O(N)\)

実装のポイント

  • 生徒番号は1-indexedなので、内部処理では0-indexedに変換(-1する)。
  • sys.stdin.read を使って高速入力処理を行い、TLEを回避。
  • 出力は最後に一度に print('\n'.join(results)) で行うことで高速化。
  • unite時にサイズ比較を行い、小さい方を大きい方に繋ぐことで効率化。
## ソースコード

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

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

    def unite(self, x, y):
        px, py = self.find(x), self.find(y)
        if px == py:
            return
        if self.size[px] < self.size[py]:
            px, py = py, px
        self.parent[py] = px
        self.size[px] += self.size[py]

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

import sys
input = sys.stdin.read

def main():
    data = input().split()
    idx = 0
    N = int(data[idx]); idx += 1
    Q = int(data[idx]); idx += 1

    uf = UnionFind(N)

    results = []
    for _ in range(Q):
        t = int(data[idx]); idx += 1
        if t == 1:
            a = int(data[idx]) - 1; idx += 1
            b = int(data[idx]) - 1; idx += 1
            uf.unite(a, b)
        else:
            x = int(data[idx]) - 1; idx += 1
            results.append(str(uf.get_size(x)))

    print('\n'.join(results))

if __name__ == "__main__":
    main()

この解説は qwen3-coder-480b によって生成されました。

投稿日時:
最終更新: