トップ > Python > UnionFind木

UnionFind木

Python UnionFind グラフ グラフ-木 AtCoder

実装

基本形

p = [-1] * (N + 1)  # 親の配列 (いなければ-1), N要素 (1-indexed, P[0]はダミー)
def root(x):
  # 親がいなければ自分(=根)を返す
  if p[x] < 0:
    return x
  # 自分の根 = 親の根
  #  & パス圧縮 (見つけた根を覚えておく)
  p[x] = root(p[x])
  return p[x]
 
def unite(x, y):
  x = root(x)
  y = root(y)
  # すでにマージされている(根が同じ)なら何もしない
  if x == y:
    return
  ## 次を入れるとxxのグループの要素数が次で計算できる: -p[root(xx)]
  ## p[x] += p[y]
  # マージされていないなら， yの親をxにする
  p[y] = x

## グループの要素数を計算する場合
## (unite でp[x] += p[y]が必要)
#def size(x):
#  x = root(x)
#  return -p[x]

コメント無し版：

# --- UnionFind ---
p = [-1] * (N + 1)
def root(x):
    if p[x] < 0:
        return x
    p[x] = root(p[x])
    return p[x]
 
def unite(x, y):
    x = root(x)
    y = root(y)
    if x == y:
        return
    p[x] += p[y]
    p[y] = x

def size(x):
    x = root(x)
    return -p[x]
# --- UnionFind ---


例題

AtCoder Typical Contest 001

上のコードそのままでACできる（Find一回あたりの均し計算量$O(Q \log N)$）
B - Union Find

N, Q = map(int, input().split())

p = [-1] * (N + 1)
def root(x):
  if p[x] < 0:
    return x
  p[x] = root(p[x])
  return p[x]
 
def unite(x, y):
  x = root(x)
  y = root(y)
  if x == y:
    return
  p[y] = x
  
for _ in range(Q):
  P, A, B = map(int, input().split())
  if P == 0:
    unite(A, B)
  else:
    print('Yes' if root(A) == root(B) else 'No')

参考