Inspired by this question, I decided to implement a Disjoint-set data structure in Python 3. I mainly followed this description for understanding the algorithm (but did not do the optimizations for path compression, union by rank etc. -- so mine is a "naive" implementation). Here is the Wikipedia article for reference.
In short, the data structure can hold n disjoint sets, and do two operations on them:
- merge any two sets into a single one
- for any element, tell to which set it belongs
My implementation is as follows:
#!/usr/bin/python class Disjoint: def __init__(self): self.sets =  def createSet(self, repr): self.sets.append([repr]) def mergeSets(self, repr1, repr2): set1 = self.findSet(repr1); set2 = self.findSet(repr2); if set1 != set2: set1.extend(set2); self.sets.remove(set2); def findSet(self, repr1): for oneSet in self.sets: if repr1 in oneSet: return oneSet def getSets(self): return self.sets;
The test class:
import unittest import disjoint class TestSequenceFunctions(unittest.TestCase): def setUp(self): pass; def test_empty(self): dis = disjoint.Disjoint(); self.assertEqual(, dis.getSets()) self.assertEqual(None, dis.findSet(1)) def test_init(self): dis = disjoint.Disjoint(); for i in range(1, 6): dis.createSet(i); for i in range(1, 6): found = dis.findSet(i); self.assertEqual(1, len(found)) self.assertEqual(i, found) expected = [[i] for i in range(1, 6)] self.assertEqual(expected, dis.getSets()) def test_simple(self): dis = disjoint.Disjoint(); for i in range(1, 6): dis.createSet(i); pairs = [[1, 2], [2, 4], [4, 5]] for p in pairs: p1 = p; p2 = p; if dis.findSet(p1) != dis.findSet(p2): dis.mergeSets(p1, p2); expetctedSets = [[1, 2, 4, 5], ]; self.assertEqual(expetctedSets, dis.getSets()) if __name__ == '__main__': unittest.main()
Any suggestion for improvement is welcome, but in particular, I'm interested in the following:
- Do you see any problem in the (naive) implementation of the algorithm?
- Can the code be improved in any way to have similar functionality, but either be more elegant or more efficient in Python? (I'm thinking of things like using list comprehensions instead of for-itertions, etc.)
- Is there any other test case that would be worth adding?
N.B.: For now, I'm not interested in how to implement the optimizations suggested in the linked source (that will be the next step, also using any new unit tests suggested here).