maintain tree

pydatastructs/miscellaneous_data_structures/disjoint_set.py

@@ -86,9 +86,20 @@ def make_root(self, key):
         if self.tree.get(key, None) is None:
             raise KeyError("Invalid key, %s"%(key))
 
-        current_root = self.find_root(key)
-        if current_root.key != key:
-            key_set = self.tree[key]
-            current_root.parent = key_set
-            key_set.size = current_root.size
+        key_set = self.tree[key]
+        if key_set.parent is not key_set:
+            current_parent = key_set.parent
+            # Remove this key subtree size from all its ancestors
+            while current_parent.parent is not current_parent:
+                current_parent.size -= key_set.size
+                current_parent = current_parent.parent
+
+            all_set_size = current_parent.size # This is the root node
+            current_parent.size -= key_set.size
+
+            # Make parent of current root as key
+            current_parent.parent = key_set
+            # size of new root will be same as previous root's size
+            key_set.size = all_set_size
+            # Make parent of key as itself
             key_set.parent = key_set

pydatastructs/miscellaneous_data_structures/tests/test_disjoint_set.py

@@ -29,3 +29,36 @@ def test_DisjointSetForest():
     assert dst.find_root(1).key == 5
     assert dst.find_root(5).key == 5
     assert raises(KeyError, lambda: dst.make_root(9))
+
+    dst = DisjointSetForest()
+    for i in range(6):
+        dst.make_set(i)
+    assert dst.tree[2].size == 1
+    dst.union(2, 3)
+    assert dst.tree[2].size == 2
+    assert dst.tree[3].size == 1
+    dst.union(1, 4)
+    dst.union(2, 4)
+    # current tree
+    ###############
+    #       2
+    #      / \
+    #     1   3
+    #    /
+    #   4
+    ###############
+    assert dst.tree[2].size == 4
+    assert dst.tree[1].size == 2
+    assert dst.tree[3].size == dst.tree[4].size == 1
+    dst.make_root(4)
+    # New tree
+    ###############
+    #       4
+    #       |
+    #       2
+    #      / \
+    #     1   3
+    ###############
+    assert dst.tree[4].size == 4
+    assert dst.tree[2].size == 3
+    assert dst.tree[1].size == dst.tree[3].size == 1