uf.py

# union find data structure
# general for any key (not just integer elements)
# $ py uf.py in/tinyUF.txt 
import sys
from collections import defaultdict

# weighted union-find data structure with path compression
class UF:
    def __init__(self):
        self._sz = defaultdict(int)
        self._id = {}

    def _root(self, x):
        # if not found, initialize element as it's own set
        if x not in self._id:
            self._id[x] = x
        while x != self._id[x]:
            self._id[x] = self._id[self._id[x]] # update to grandparent
            x = self._id[x]
        return x

    def find(self, p, q):
        return self._root(p) == self._root(q)

    def union(self, p, q):
        x = self._root(p)
        y = self._root(q)
        if x == y:
            return
        if self._sz[x] < self._sz[y]:
            self._id[x] = y
            self._sz[y] += self._sz[x]
        else:
            self._id[y] = x
            self._sz[x] += self._sz[y]

    def __str__(self):
        return str(self._id)

def main():
    uf = UF()
    filename = sys.argv[1]
    with open(filename, 'r') as f:
        line = f.readline().strip()
        while line:
            v, w = line.strip().split()
            uf.union(v, w)
            line = f.readline().strip()
    print(uf)

if __name__ == '__main__':
    main()