Implement find_size for disjoint set union data structure (#516)

Author: sak-codes

pydatastructs/miscellaneous_data_structures/disjoint_set.py

@@ -115,3 +115,12 @@ def make_root(self, key):
             key_set.size = all_set_size
             # Make parent of key as itself
             key_set.parent = key_set
+
+    def find_size(self, key):
+        """
+        Finds the size of set to which the key belongs.
+        """
+        if self.tree.get(key, None) is None:
+            raise KeyError("Invalid key, %s"%(key))
+
+        return self.find_root(key).size

pydatastructs/miscellaneous_data_structures/tests/test_disjoint_set.py

@@ -9,16 +9,19 @@ def test_DisjointSetForest():
 
     dst.union(1, 2)
     dst.union(1, 5)
+    assert dst.find_size(2) == 3
     dst.union(1, 6)
     dst.union(1, 8)
     dst.union(3, 4)
+    assert dst.find_size(3) == 2
 
     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(3) == dst.find_root(4)
     assert dst.find_root(7).key == 7
 
     assert raises(KeyError, lambda: dst.find_root(9))
+    assert raises(KeyError, lambda: dst.find_size(9))
     dst.union(3, 1)
     assert dst.find_root(3).key == 1
     assert dst.find_root(5).key == 1