disjointset.py

sets = [[i] for i in range(len(nums))]
parent = [i for i in range(len(nums))]

def union(a, b):
  set_a = parent[a]
  set_b = parent[b]

  if set_a == set_b:
    return

  if len(sets[set_a]) < len(sets[set_b]):
    for a in sets[set_a]:
      sets[set_b].append(a)
      parent[a] = parent[b]

    sets[set_a] = None
  else:
    for b in sets[set_b]:
      sets[set_a].append(b)
      parent[b] = parent[a]

        sets[set_b] = None

disjointset_size.py

parent = [i for i in range(len(nums))]
size = [0] * len(nums)

def union(a, b):
  set_a = parent[a]
  set_b = parent[b]
  
  if set_a == set_b:
    return
  
  if size[set_a] < size[set_b]:
    parent[set_a] = set_b
    size[set_b] += set_a
  else:
    parent[set_b] = set_a
    size[set_a] = set_b

def find_parent(a):
  p = parent[a]
  
  if p == a:
    return a
  
  parent[a] = find_parent(p)
  
  return parent[a]