add make_root method

Author: Smit-create

pydatastructs/miscellaneous_data_structures/disjoint_set.py

@@ -16,6 +16,9 @@ class DisjointSetForest(object):
     >>> dst.union(1, 2)
     >>> dst.find_root(2).key
     1
+    >>> dst.make_root(2)
+    >>> dst.find_root(2).key
+    2
 
     References
     ==========
@@ -74,3 +77,18 @@ def union(self, key1, key2):
 
             y_root.parent = x_root
             x_root.size += y_root.size
+
+    def make_root(self, key):
+        """
+        Finds the set to which the key belongs
+        and makes it as the root of the set.
+        """
+        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.parent = key_set

pydatastructs/miscellaneous_data_structures/tests/test_disjoint_set.py

@@ -15,9 +15,15 @@ def test_DisjointSetForest():
 
     assert (dst.find_root(1) == dst.find_root(2) ==
             dst.find_root(5) == dst.find_root(6) == dst.find_root(8))
+    assert dst.find_root(1).key == 1
+    assert dst.find_root(2).key == 1
+    dst.make_root(8)
+    assert dst.find_root(2).key == 8
+    assert dst.find_root(8).key == 8
     assert dst.find_root(3) == dst.find_root(4)
     assert dst.find_root(7).key == 7
-
+    dst.make_root(2)
     assert raises(KeyError, lambda: dst.find_root(9))
+    assert raises(KeyError, lambda: dst.make_root(19))
     dst.union(3, 1)
-    assert dst.find_root(3).key == 1
+    assert dst.find_root(3).key == 2