Added Disjoint set data structure (#143)

Author: czgdp1807

pydatastructs/miscellaneous_data_structures/__init__.py

@@ -2,7 +2,9 @@
 
 from . import (
     stack,
-    binomial_trees
+    binomial_trees,
+    queue,
+    disjoint_set
 )
 
 from .binomial_trees import (
@@ -19,3 +21,8 @@
     Queue,
 )
 __all__.extend(queue.__all__)
+
+from .disjoint_set import (
+    DisjointSetForest,
+)
+__all__.extend(disjoint_set.__all__)

pydatastructs/miscellaneous_data_structures/disjoint_set.py

@@ -0,0 +1,72 @@
+from pydatastructs.utils import Set
+
+__all__ = ['DisjointSetForest']
+
+class DisjointSetForest(object):
+    """
+    Represents a forest of disjoint set trees.
+
+    Examples
+    ========
+
+    >>> from pydatastructs import DisjointSetForest
+    >>> dst = DisjointSetForest()
+    >>> dst.make_set(1)
+    >>> dst.make_set(2)
+    >>> dst.union(1, 2)
+    >>> dst.find_root(2).key
+    1
+
+    References
+    ==========
+
+    .. [1] https://en.wikipedia.org/wiki/Disjoint-set_data_structure
+    """
+
+    __slots__ = ['tree']
+
+    def __new__(cls):
+        obj = object.__new__(cls)
+        obj.tree = dict()
+        return obj
+
+    def make_set(self, key, data=None):
+        """
+        Adds a singleton set to the tree
+        of disjoint sets with given key
+        and optionally data.
+        """
+        if self.tree.get(key, None) is None:
+            new_set = Set(key, data)
+            self.tree[key] = new_set
+            new_set.parent = new_set
+            new_set.size = 1
+
+    def find_root(self, key):
+        """
+        Finds the root of the set
+        with the given key by path
+        splitting algorithm.
+        """
+        if self.tree.get(key, None) is None:
+            raise KeyError("Invalid key, %s"%(key))
+        _set = self.tree[key]
+        while _set.parent is not _set:
+            _set, _set.parent = _set.parent, _set.parent.parent
+        return _set
+
+    def union(self, key1, key2):
+        """
+        Takes the union of the two
+        disjoint set trees with given
+        keys. The union is done by size.
+        """
+        x_root = self.find_root(key1)
+        y_root = self.find_root(key2)
+
+        if x_root is not y_root:
+            if x_root.size < y_root.size:
+                x_root, y_root = y_root, x_root
+
+            y_root.parent = x_root
+            x_root.size += y_root.size

pydatastructs/miscellaneous_data_structures/tests/test_disjoint_set.py

@@ -0,0 +1,23 @@
+from pydatastructs import DisjointSetForest
+from pydatastructs.utils.raises_util import raises
+
+def test_DisjointSetForest():
+
+    dst = DisjointSetForest()
+    for i in range(8):
+        dst.make_set(i+1)
+
+    dst.union(1, 2)
+    dst.union(1, 5)
+    dst.union(1, 6)
+    dst.union(1, 8)
+    dst.union(3, 4)
+
+    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))
+    dst.union(3, 1)
+    assert dst.find_root(3).key == 1

pydatastructs/utils/__init__.py

@@ -7,6 +7,7 @@
     BinomialTreeNode,
     AdjacencyListGraphNode,
     AdjacencyMatrixGraphNode,
-    GraphEdge
+    GraphEdge,
+    Set
 )
 __all__.extend(misc_util.__all__)

pydatastructs/utils/misc_util.py

@@ -4,7 +4,8 @@
     'BinomialTreeNode',
     'AdjacencyListGraphNode',
     'AdjacencyMatrixGraphNode',
-    'GraphEdge'
+    'GraphEdge',
+    'Set'
 ]
 
 _check_type = lambda a, t: isinstance(a, t)
@@ -225,3 +226,26 @@ def __new__(cls, node1, node2, value=None):
 
     def __str__(self):
         return str((self.source.name, self.target.name))
+
+class Set(object):
+    """
+    Represents a set in a forest of disjoint sets.
+
+    Parameters
+    ==========
+
+    key: Hashable python object
+        The key which uniquely identifies
+        the set.
+    data: Python object
+        The data to be stored in the set.
+    """
+
+    __slots__ = ['parent', 'size', 'key', 'data']
+
+    def __new__(cls, key, data):
+        obj = object.__new__(cls)
+        obj.key = key
+        obj.data = data
+        obj.parent, obj.size = [None]*2
+        return obj