Union-Find data structure implementation with examples (Kruskala MST, Tarjan LCA etc)

uf.py

# Union-Find (Disjoint Set Union)
# more information on wiki:
# http://en.wikipedia.org/wiki/Disjoint-set_data_structure

class UF(object):

    def __init__(self, size):
        self.p = [None]*size
        self.rank = [1]*size

    def make(self, el):
        self.p[el] = el 

    def find(self, el):
        if self.p[el] == el: return el 
        self.p[el] = self.find(self.p[el])
        return self.p[el]

    def unite(self, l, r):
        li, ri = self.find(l), self.find(r)
        if self.rank[li] < self.rank[ri]:
            self.p[li] = ri
        else:
            self.p[ri] = li 
            if self.rank[li] == self.rank[ri]:
                self.rank[li] += 1

print ""
print "== compontents =="

g = UF(10)
map(g.make, range(10))
print g.p
for l,r in [(0,1), (2,3), (4,5), (6,7), (7,8), (5,6), (6,9)]:
    g.unite(l,r)

print "num", sum(1 for l,r in enumerate(g.p) if l==r)

print ""
print "== Kruskala =="

from operator import itemgetter

edges = [(0,1,1),
         (1,2,100),
         (2,5,40), 
         (3,4,60), 
         (4,5,12), 
         (5,6,20), 
         (7,8,40), 
         (8,9,80), 
         (6,7,21), 
         (7,10,52)]

krs = UF(11)
map(krs.make, range(11))

path = []
for l,r,_ in sorted(edges, key=itemgetter(2)):
    if krs.find(l) != krs.find(r):
        path.append((l,r))
        krs.unite(l,r)

print path

print ""
print "=== Tarjan LCA ==="
# tree
# 0 -> 1
#   -> 2
#   -> 3 -> 5
#   -> 4
tree = {
    0: [1, 2, 3, 4],
    2: [],
    3: [5],
    4: [],
    5: []
}

# list of queries
queries = [
    (3, 5), # 3
    (1, 2), # 0
    (1, 5)  # 0
]


from collections import defaultdict

q = defaultdict(set)
for f,t in queries:
    q[f].add(t)
    q[t].add(f)


seen, ancs, t = [False]*6, [None]*6, UF(6)
def dfs(node):
    ancs[node] = node
    t.make(node)

    for k in tree.get(node, ()):
        if not seen[k]:
            dfs(k)
            ancs[k] = node 
            t.unite(k, node)

    for qt in q.get(node, ()):
        if seen[qt]:
            print node, "&", qt, "-->", ancs[t.find(qt)]

    seen[node] = True

dfs(0)