Given lists of related numbers, merge related lists to create disjoint sets

Question

Given:

[(1,2),(3,4),(5,6),(3,7),(5,7)]

Output:

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

Explanation:

(1,2)
(1,2), (3,4)
(1,2), (3,4), (5,6)
(1,2), (3,4,7), (5,6)
(1,2), (3,4,7,5,6)

I have written a lousy algorithm:

Case 1: both numbers in pair are new (never seen before):
    Make a new set with these two numbers
Case 2: one of the number in pair is new, other is already a part of some set:
    Merge the new number in other's set
Case 3: both the numbers belong to some set:
    Union the second set into first. Destroy the second set.

Is there a more pythonic (fancy) solution to this algo?

Answer 1

You can use the Unionfind algorithm for this. First, we are using a dictionary to create a tree from the pairs:

leaders = collections.defaultdict(lambda: None)

Now we use two functions -- union and find -- to populate that tree:

def find(x):
    l = leaders[x]
    if l is not None:
        l = find(l)
        leaders[x] = l
        return l
    return x

def union(x, y):
    lx, ly = find(x), find(y)
    if lx != ly:
        leaders[lx] = ly

Just iterate over all the pairs and put them into the tree.

for a, b in [(1,2),(3,4),(5,6),(3,7),(5,7)]:
    union(a, b)

It then looks like this: {1: 2, 2: None, 3: 4, 4: 7, 5: 6, 6: 7, 7: None}

Finally, we group the numbers by their respective "leaders", i.e. what is returned by find:

groups = collections.defaultdict(set)
for x in leaders:
    groups[find(x)].add(x)

Now, groups.values() is [set([1, 2]), set([3, 4, 5, 6, 7])]

Complexity should be something on the order of O(nlogn).