Union find implementation using Python

Solution that runs in O(n) time

def indices_dict(lis):
    d = defaultdict(list)
    for i,(a,b) in enumerate(lis):
        d[a].append(i)
        d[b].append(i)
    return d

def disjoint_indices(lis):
    d = indices_dict(lis)
    sets = []
    while len(d):
        que = set(d.popitem()[1])
        ind = set()
        while len(que):
            ind |= que 
            que = set([y for i in que 
                         for x in lis[i] 
                         for y in d.pop(x, [])]) - ind
        sets += [ind]
    return sets

def disjoint_sets(lis):
    return [set([x for i in s for x in lis[i]]) for s in disjoint_indices(lis)]

---

How it works:

>>> lis = [(1,2),(2,3),(4,5),(6,7),(1,7)]
>>> indices_dict(lis)
>>> {1: [0, 4], 2: [0, 1], 3: [1], 4: [2], 5: [2], 6: [3], 7: [3, 4]})

indices_dict gives a map from an equivalence # to an index in lis. E.g. 1 is mapped to index 0 and 4 in lis.

>>> disjoint_indices(lis)
>>> [set([0,1,3,4], set([2])]

disjoint_indices gives a list of disjoint sets of indices. Each set corresponds to indices in an equivalence. E.g. lis[0] and lis[3] are in the same equivalence but not lis[2].

>>> disjoint_set(lis)
>>> [set([1, 2, 3, 6, 7]), set([4, 5])]

disjoint_set converts disjoint indices into into their proper equivalences.

---

Time complexity

The O(n) time complexity is difficult to see but I’ll try to explain. Here I will use n = len(lis).

1. indices_dict certainly runs in O(n) time because only 1 for-loop

2. disjoint_indices is the hardest to see. It certainly runs in O(len(d)) time since the outer loop stops when d is empty and the inner loop removes an element of d each iteration. now, the len(d) <= 2n since d is a map from equivalence #’s to indices and there are at most 2n different equivalence #’s in lis. Therefore, the function runs in O(n).

3. disjoint_sets is difficult to see because of the 3 combined for-loops. However, you’ll notice that at most i can run over all n indices in lis and x runs over the 2-tuple, so the total complexity is 2n = O(n)