disjoint_set example for https://www.v2ex.com/t/799577

disjoint_set_example.py

input = [
    {'A': 'A1'},
    {'A1': 'A2'},
    {'A3': 'A4'},
    {'A2': 'A3'},
    {'B4': 'B3'},
    {'B1': 'B4'},
    {'B4': 'B'},
    {'B3': 'B2'},
    {'B2': 'B5'},
]

dset = {}
output = {}

def find(k):
  if k not in dset:
    dset[k] = k
  elif dset[k] != k:
    dset[k] = find(dset[k])
  return dset[k]

for pairs in input:
  for (k, v) in pairs.items():
    vp = find(v)
    dset[vp] = find(k)

print(dset)

for pairs in input:
  for (k, v) in pairs.items():
    kp = find(k)
    if kp not in output:
      output[kp] = set()
    output[kp].add(k)
    output[kp].add(v)

print(output)