UnionFind data structure - (by rank) (Inspired from networkx.utils.unionfind)

union_find_by_rank.py

"""
Inspired from the great Graph lib.: networkx.utils.unionfind

The one implemented in networkx was union by weight and implicitly adds element upon __getitem__. I changed it to union-by-rank and 
have a add() call equivalent to the make_set() call in the original data structure.
to_sets() changed a bit too to remove dependency to networkx.

Other than above most of the code is same.
"""

class UnionFindByRank:
    def __init__(self, *elements):
        self.parents = {}
        self.ranks = {}

        self.add(*elements)

    def add(self, *elements):
        for x in elements:
            self.ranks[x] = 0
            self.parents[x] = x

    def __getitem__(self, object):
        # find path of objects leading to the root
        path = [object]
        root = self.parents[object]
        while root != path[-1]:
            path.append(root)
            root = self.parents[root]

        # compress the path and return
        for ancestor in path:
            self.parents[ancestor] = root
        return root

    def __iter__(self):
        return iter(self.parents)

    def to_sets(self):
        import collections
        d = collections.defaultdict(set)
        for elem in self.parents.keys():
            d[self[elem]].add(elem)
        return d.values()

    def union(self, *objects):
        roots = [self[x] for x in objects]
        tallest = max(roots, key=lambda r: self.ranks[r])

        for x in objects:
            root = self[x]
            if root != tallest:
                if self.ranks[root] == self.ranks[tallest]:
                    self.ranks[tallest] += 1
                self.parents[root] = tallest