Implement disjoint_sets method in DisjointSetForest (#521)

Author: sak-codes

pydatastructs/miscellaneous_data_structures/disjoint_set.py

@@ -124,3 +124,20 @@ def find_size(self, key):
             raise KeyError("Invalid key, %s"%(key))
 
         return self.find_root(key).size
+
+    def disjoint_sets(self):
+        """
+        Returns a list of disjoint sets in the data structure.
+        """
+        result = dict()
+        for key in self.tree.keys():
+            parent = self.find_root(key).key
+            members = result.get(parent, [])
+            members.append(key)
+            result[parent] = members
+        sorted_groups = []
+        for v in result.values():
+            sorted_groups.append(v)
+            sorted_groups[-1].sort()
+        sorted_groups.sort()
+        return sorted_groups

pydatastructs/miscellaneous_data_structures/tests/test_disjoint_set.py

@@ -17,6 +17,7 @@ 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.disjoint_sets() == [[1, 2, 5, 6, 8], [3, 4], [7]]
     assert dst.find_root(3) == dst.find_root(4)
     assert dst.find_root(7).key == 7
 
@@ -26,6 +27,7 @@ def test_DisjointSetForest():
     assert dst.find_root(3).key == 1
     assert dst.find_root(5).key == 1
     dst.make_root(6)
+    assert dst.disjoint_sets() == [[1, 2, 3, 4, 5, 6, 8], [7]]
     assert dst.find_root(3).key == 6
     assert dst.find_root(5).key == 6
     dst.make_root(5)
@@ -42,6 +44,7 @@ def test_DisjointSetForest():
     assert dst.tree[3].size == 1
     dst.union(1, 4)
     dst.union(2, 4)
+    assert dst.disjoint_sets() == [[0], [1, 2, 3, 4], [5]]
     # current tree
     ###############
     #       2