～ Union-Find ～ チートシート

目次

Union-Findとは
実装
工夫

はじめに

チートシートの扱いついてはここを読んでください

Union-Findとは

要素同士の結合から、グループの数や要素数などを求めるためのデータ構造
要素全体を森、同じ木に属する要素を同じグループとして扱う

Union-Find木を使う手順
①すべての要素を独立した木とし、自身を根とする
②２つの要素の結合が与えられそれぞれが異なる木の場合、２つの木を結合し１つの木とする
③根が同じ要素は同じ木に属することを利用し、要素同士が同一のグループかを判定する

詳しくはこちらのページのスライドを参照

実装

classを使うのに不安を覚える雑魚 用のコード

Union-Find.py

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

root = [i for i in range(N+1)] # 要素の根を格納する配列
rank = [0]*(N+1) # 木のランク（高さ）を格納する配列
size = [1]*(N+1) # 木のサイズ（要素数）を格納する配列

def union(x,y):
  root_x = find(x)
  root_y = find(y)
  if root_x != root_y:
    if rank[root_x] >= rank[root_y]:
      if rank[root_x] == rank[root_y]:
        rank[root_x] += 1
      root[root_y] = root_x
      size[root_x] += size[root_y]
    else:
      root[root_x] = root_y
      size[root_y] += size[root_x]

def find(x):
  tmp_root = x
  while True:
    if tmp_root == root[tmp_root]:
      break
    tmp_root = root[tmp_root]
  root[x] = tmp_root
  return tmp_root

# find(A) == find(B) ならばAとBは同じグループに属する
# size[find(A)] でAを含むグループの要素数を得られる

for i in range(M):
  A,B = map(int,input().split())
  union(A,B)

工夫

計算量を落とすために、主に２つの工夫を行う

①ランクの利用

根が等しいかどうかで要素のグループが等しいかを判定するので、木同士を結合する際に片方の木の根をもう片方の木の根に直接つなぐ
この時、ランク（木の高さ）が小さくなるようにランクが小さい方の根を大きい方の根につなぐ

②経路圧縮

途中で要素の根がわかったら、その要素を根に直接つなぎなおす（要素の根を格納する配列を更新する）