RedBlackTrees search()

codezonediitj · Jun 11, 2024 · d70191e · d70191e
1 parent 10ef274
commit d70191e
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/pydatastructs/trees/_backend/cpp/RedBlackTree.hpp b/pydatastructs/trees/_backend/cpp/RedBlackTree.hpp
@@ -574,10 +574,15 @@ static PyObject* RedBlackTree_delete(RedBlackTree* self, PyObject *args, PyObjec
     Py_RETURN_NONE;
 }
 
+static PyObject* RedBlackTree_search(RedBlackTree* self, PyObject *args, PyObject *kwds) {
+    return BinarySearchTree_search(self->sbbt->bst, args, kwds);
+}
+
 
 static struct PyMethodDef RedBlackTree_PyMethodDef[] = {
     {"insert", (PyCFunction) RedBlackTree_insert, METH_VARARGS, NULL},
     {"delete", (PyCFunction) RedBlackTree_delete, METH_VARARGS | METH_KEYWORDS, NULL},
+    {"search", (PyCFunction) RedBlackTree_search, METH_VARARGS | METH_KEYWORDS, NULL},
     {"lower_bound", (PyCFunction) RedBlackTree_lower_bound, METH_VARARGS | METH_KEYWORDS, NULL},
     {"upper_bound", (PyCFunction) RedBlackTree_upper_bound, METH_VARARGS | METH_KEYWORDS, NULL},
     {"_get_parent", (PyCFunction) RedBlackTree__get_parent, METH_VARARGS, NULL},

diff --git a/pydatastructs/trees/tests/test_binary_trees.py b/pydatastructs/trees/tests/test_binary_trees.py
@@ -700,9 +700,15 @@ def _test_RedBlackTree(backend):
     assert [node.key for node in in_order] == [5, 10, 20, 30, 40, 50]
     assert [node.key for node in pre_order] == [40, 20, 10, 5, 30, 50]
 
+    assert tree.search(50) == 0
+    assert tree.search(20) == 3
+    assert tree.search(30) == 2
     tree.delete(50)
     tree.delete(20)
     tree.delete(30)
+    assert tree.search(50) is None
+    assert tree.search(20) is None
+    assert tree.search(30) is None
 
     in_order = trav.depth_first_search(order='in_order')
     pre_order = trav.depth_first_search(order='pre_order')