Added disjoint sets data structure.

jwasham · jwasham · commit 47cc6ff8841a · 2016-09-22T19:49:17.000-07:00
diff --git a/disjoint-sets/disjoint-sets.py b/disjoint-sets/disjoint-sets.py
@@ -0,0 +1,68 @@
+from array import array
+import os
+
+
+class DisjointSets(object):
+
+    def __init__(self, size):
+        self.hierarchy = array('l', [0 for _ in range(size)])
+        self.names = {}
+
+    def add(self, item_id, name):
+        self.hierarchy[item_id] = -1
+        self.names[item_id] = name
+
+    def union(self, root1, root2):
+        if self.hierarchy[root1] <= self.hierarchy[root2]:   # root1 is a larger tree, since roots are -size of tree
+            self.hierarchy[root1] += self.hierarchy[root2]   # adding to increase negative value
+            self.hierarchy[root2] = root1
+        else:
+            self.hierarchy[root2] += self.hierarchy[root1]
+            self.hierarchy[root1] = root2
+
+    def find(self, item_id):
+        '''
+        Finds the root of the set of which item_id is a member.
+        To speed up subsequent finds, updates the parent to root at each level in path to root.
+        :param item_id:
+        :return: integer representative of set
+        '''
+        if self.hierarchy[item_id] < 0:
+            return item_id
+        else:
+            # path compression
+            self.hierarchy[item_id] = self.find(self.hierarchy[item_id])
+            return self.hierarchy[item_id]
+
+    def __str__(self):
+        ret_str = repr(self.hierarchy)
+        ret_str += os.linesep + repr(self.names)
+        return ret_str
+
+
+def main():
+    ds = DisjointSets(10)
+    ds.add(0, "Microsoft")
+    ds.add(1, "WebTV")
+    ds.add(2, "Google")
+    ds.add(3, "DeepMind")
+    ds.add(4, "Skype")
+    ds.add(5, "Uber")
+    ds.add(6, "TinyCo")
+    ds.add(7, "TeenyTinyCo")
+
+    ds.union(0, 1)
+    ds.union(2, 3)
+    ds.union(0, 4)
+    ds.union(4, 6)
+    ds.union(6, 7)
+
+    # print(ds)
+
+    assert(ds.find(3) == 2)
+
+    # print(ds)
+
+
+if __name__ == '__main__':
+    main()
