UNB/ CS/ David Bremner/ teaching/ cs3383/ lectures/ 23.0-demos/ union.py
#!/usr/bin/env python3
class Partition:
    def __init__(P,n):
      # sometimes called makeset(j)
      P.parent = [j for j in range(n)]
      P.rank = [0] * n

    def find(P, key):
      while P.parent[key] != key:
        key = P.parent[key]
      return key

    def union(P,x,y):
      rx = P.find(x)
      ry = P.find(y)
      if rx != ry:
        if P.rank[rx] > P.rank[ry]:
          P.parent[ry] = rx
        else:
          P.parent[rx] = ry
          if P.rank[rx] == P.rank[ry]:
            P.rank[ry] += 1

    def print(P,name):
        with open("{:s}.dot".format(name),'w') as f:
            print('digraph "{:s}" {{ rankdir=\"BT\"'.format(name),file=f)
            for i in range(len(P.parent)):
                if P.parent[i] == i:
                    print("{:d}".format(i), file=f)
                else:
                    print("{:d} -> {:d}".format(i,P.parent[i]),file=f)
            print("}",file=f)

if __name__ == "__main__":
    P=Partition(7)
    P.union(0,3)
    P.union(1,4)
    P.union(2,5)
    P.print("step1")
    P.union(2,6)
    P.union(4,0)
    P.print("step2")
    P.union(1,6)
    P.print("step3")