A simple (but complete) union-find structure in Python.

union_find.py

# Author: Guillermo Angeris (@guillean)
# Simple UnionFind based on https://www.cs.princeton.edu/~rs/AlgsDS07/01UnionFind.pdf

class UnionFind:
    def __init__(self):
        self.rank = {}
        self.parent = {}

    def add(self, elem):
        self.rank[elem] = 1
        self.parent[elem] = elem

    def find(self, a):
        while a != self.parent[a]:
            self.parent[a] = self.parent[self.parent[a]]
            a = self.parent[a]
        return a

    def union(self, a, b):
        i_a, i_b = self.find(a), self.find(b)
        if i_a == i_b: return
        if self.rank[a] < self.rank[b]:
            self.parent[i_a] = i_b
            self.rank[i_a] += self.rank[i_b]
        else:
            self.parent[i_b] = i_a
            self.rank[i_b] += self.rank[i_a]

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

    def __contains__(self, a):
        return a in self.parent