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)
.
indices_dict
certainly runs inO(n)
time because only 1 for-loopdisjoint_indices
is the hardest to see. It certainly runs inO(len(d))
time since the outer loop stops whend
is empty and the inner loop removes an element ofd
each iteration. now, thelen(d) <= 2n
sinced
is a map from equivalence #'s to indices and there are at most2n
different equivalence #'s inlis
. Therefore, the function runs inO(n)
.disjoint_sets
is difficult to see because of the 3 combined for-loops. However, you'll notice that at mosti
can run over alln
indices inlis
andx
runs over the 2-tuple, so the total complexity is2n = O(n)
I think this is an elegant solution, using the built in set functions:
#!/usr/bin/python3
def union_find(lis):
lis = map(set, lis)
unions = []
for item in lis:
temp = []
for s in unions:
if not s.isdisjoint(item):
item = s.union(item)
else:
temp.append(s)
temp.append(item)
unions = temp
return unions
if __name__ == '__main__':
l = [[1, 2], [2, 3], [4, 5], [6, 7], [1, 7]]
print(union_find(l))
It returns a list of sets.