union_find.py

parent = [-1] * n

def find(parent, i):
  if parent[i] < 0:
    return i
  parent[i] = find(parent, parent[i])
  return parent[i]

def unite(parent, i, j):
  i = find(parent, i)
  j = find(parent, j)
  if i == j:
    return
  parent[j] += parent[i]
  parent[i] = j

def size(parent, i):
  return -parent[find(parent, i)]