Add upper and lower bound in BST (#537)

Author: sak-codes

pydatastructs/trees/binary_trees.py

@@ -275,6 +275,98 @@ def search(self, key, **kwargs):
                 walk = self.tree[walk].right
         return (walk, parent) if ret_parent else walk
 
+    def _bound_helper(self, node_idx, bound_key, is_upper=False):
+        if node_idx is None:
+            return None
+        if self.tree[node_idx].key is None:
+            return None
+
+        if self.tree[node_idx].key == bound_key:
+            if not is_upper:
+                return self.tree[node_idx].key
+            else:
+                return self._bound_helper(self.tree[node_idx].right,
+                                            bound_key, is_upper)
+
+        if self.comparator(self.tree[node_idx].key, bound_key):
+            return self._bound_helper(self.tree[node_idx].right,
+                                            bound_key, is_upper)
+        else:
+            res_bound = self._bound_helper(self.tree[node_idx].left,
+                                                   bound_key, is_upper)
+            return res_bound if res_bound is not None else self.tree[node_idx].key
+
+
+    def lower_bound(self, key, **kwargs):
+        """
+        Finds the lower bound of the given key in the tree
+
+        Parameters
+        ==========
+
+        key
+            The key for comparison
+
+        Examples
+        ========
+
+        >>> from pydatastructs.trees import BinarySearchTree as BST
+        >>> b = BST()
+        >>> b.insert(10, 10)
+        >>> b.insert(18, 18)
+        >>> b.insert(7, 7)
+        >>> b.lower_bound(9)
+        10
+        >>> b.lower_bound(7)
+        7
+        >>> b.lower_bound(20) is None
+        True
+
+        Returns
+        =======
+
+        value
+            The lower bound of the given key.
+            Returns None if the value doesn't exist
+        """
+        return self._bound_helper(self.root_idx, key)
+
+
+    def upper_bound(self, key, **kwargs):
+        """
+        Finds the upper bound of the given key in the tree
+
+        Parameters
+        ==========
+
+        key
+            The key for comparison
+
+        Examples
+        ========
+
+        >>> from pydatastructs.trees import BinarySearchTree as BST
+        >>> b = BST()
+        >>> b.insert(10, 10)
+        >>> b.insert(18, 18)
+        >>> b.insert(7, 7)
+        >>> b.upper_bound(9)
+        10
+        >>> b.upper_bound(7)
+        10
+        >>> b.upper_bound(20) is None
+        True
+
+        Returns
+        =======
+
+        value
+            The upper bound of the given key.
+            Returns None if the value doesn't exist
+        """
+        return self._bound_helper(self.root_idx, key, True)
+
+
     def delete(self, key, **kwargs):
         (walk, parent) = self.search(key, parent=True)
         a = None

pydatastructs/trees/tests/test_binary_trees.py

@@ -525,7 +525,31 @@ def test_RedBlackTree():
     assert [node.key for node in in_order] == [2, 6, 7, 10, 15, 16, 17, 18, 25, 30, 40, 60]
     assert [node.key for node in pre_order] == [16, 10, 6, 2, 7, 15, 25, 18, 17, 40, 30, 60]
 
+    assert tree.lower_bound(0) == 2
+    assert tree.lower_bound(2) == 2
+    assert tree.lower_bound(3) == 6
+    assert tree.lower_bound(7) == 7
+    assert tree.lower_bound(25) == 25
+    assert tree.lower_bound(32) == 40
+    assert tree.lower_bound(41) == 60
+    assert tree.lower_bound(60) == 60
+    assert tree.lower_bound(61) is None
+
+    assert tree.upper_bound(0) == 2
+    assert tree.upper_bound(2) == 6
+    assert tree.upper_bound(3) == 6
+    assert tree.upper_bound(7) == 10
+    assert tree.upper_bound(25) == 30
+    assert tree.upper_bound(32) == 40
+    assert tree.upper_bound(41) == 60
+    assert tree.upper_bound(60) is None
+    assert tree.upper_bound(61) is None
+
     tree = RedBlackTree()
+
+    assert tree.lower_bound(1) is None
+    assert tree.upper_bound(0) is None
+
     tree.insert(10)
     tree.insert(20)
     tree.insert(30)