select(), rank(), and set_tree()

Author: Kishan-Ved

pydatastructs/trees/_backend/cpp/AVLTree.hpp

@@ -40,6 +40,13 @@ static PyObject* AVLTree___str__(AVLTree *self) {
     return BinarySearchTree___str__(self->sbbt->bst);
 }
 
+static PyObject* AVLTree_set_tree(AVLTree* self, PyObject *args) {
+    ArrayForTrees* arr = reinterpret_cast<ArrayForTrees*>(PyObject_GetItem(args, PyZero));
+    self->sbbt->bst->binary_tree->tree = arr;
+    self->tree = self->sbbt->bst->binary_tree->tree;
+    Py_RETURN_NONE;
+}
+
 static PyObject* AVLTree_search(AVLTree* self, PyObject *args, PyObject *kwds) {
     return BinarySearchTree_search(self->sbbt->bst, args, kwds);
 }
@@ -236,10 +243,21 @@ static PyObject* AVLTree_insert(AVLTree* self, PyObject *args) {
     Py_RETURN_NONE;
 }
 
+static PyObject* AVLTree_rank(AVLTree* self, PyObject *args) {
+    return BinarySearchTree_rank(self->sbbt->bst, args);
+}
+
+static PyObject* AVLTree_select(AVLTree* self, PyObject *args) {
+    return BinarySearchTree_select(self->sbbt->bst, args);
+}
+
 static struct PyMethodDef AVLTree_PyMethodDef[] = {
     {"search", (PyCFunction) AVLTree_search, METH_VARARGS | METH_KEYWORDS, NULL},
     {"insert", (PyCFunction) AVLTree_insert, METH_VARARGS, NULL},
+    {"set_tree", (PyCFunction) AVLTree_set_tree, METH_VARARGS, NULL},
     {"balance_factor", (PyCFunction) AVLTree_balance_factor, METH_VARARGS, NULL},
+    {"rank", (PyCFunction) AVLTree_rank, METH_VARARGS, NULL},
+    {"select", (PyCFunction) AVLTree_select, METH_VARARGS, NULL},
     {"_left_right_rotate", (PyCFunction) AVLTree__left_right_rotate, METH_VARARGS, NULL},
     {"_right_left_rotate", (PyCFunction) AVLTree__right_left_rotate, METH_VARARGS, NULL},
     {"_left_rotate", (PyCFunction) AVLTree__left_rotate, METH_VARARGS, NULL},

pydatastructs/trees/_backend/cpp/BinarySearchTree.hpp

@@ -84,8 +84,9 @@ static PyObject* BinarySearchTree_search(BinarySearchTree* self, PyObject* args,
         return NULL;
     }
     BinaryTree* bt = self->binary_tree;
+    Py_INCREF(Py_None);
     PyObject* parent = Py_None;
-    PyObject* walk = PyLong_FromLong(PyLong_AsLong(bt->root_idx));
+    PyObject* walk = bt->root_idx;
 
     if (reinterpret_cast<TreeNode*>(bt->tree->_one_dimensional_array->_data[PyLong_AsLong(walk)])->key == Py_None) {
         Py_RETURN_NONE;

pydatastructs/trees/binary_trees.py

@@ -932,7 +932,7 @@ def __new__(cls, key=None, root_data=None, comp=None,
 
     @classmethod
     def methods(cls):
-        return ['__new__', 'insert', 'delete']
+        return ['__new__', 'set_tree', 'insert', 'delete']
 
     left_height = lambda self, node: self.tree[node.left].height \
                                         if node.left is not None else -1
@@ -941,6 +941,9 @@ def methods(cls):
     balance_factor = lambda self, node: self.right_height(node) - \
                                         self.left_height(node)
 
+    def set_tree(self, arr):
+        self.tree = arr
+
     def _right_rotate(self, j, k):
         super(AVLTree, self)._right_rotate(j, k)
         self.tree[j].height = max(self.left_height(self.tree[j]),

pydatastructs/trees/tests/test_binary_trees.py

@@ -214,8 +214,9 @@ def _test_AVLTree(backend):
     assert [node.key for node in in_order] == [1]
     assert [node.key for node in pre_order] == [1]
 
-    a3 = AVLTree(0,0,backend=backend)
-    for i in range(1,7):
+    a3 = AVLTree()
+    a3.set_tree( ArrayForTrees(TreeNode, 0, backend=backend) )
+    for i in range(0,7):
         a3.tree.append(TreeNode(i, i, backend=backend))
     a3.tree[0].left = 1
     a3.tree[0].right = 6
@@ -224,15 +225,17 @@ def _test_AVLTree(backend):
     a3.tree[2].left = 3
     a3.tree[2].right = 4
     a3._left_right_rotate(0, 1)
+    assert str(a3) == "[(4, 0, 0, 6), (5, 1, 1, 3), (1, 2, 2, 0), (None, 3, 3, None), (None, 4, 4, None), (None, 5, 5, None), (None, 6, 6, None)]"
 
     trav = BinaryTreeTraversal(a3, backend=backend)
     in_order = trav.depth_first_search(order='in_order')
     pre_order = trav.depth_first_search(order='pre_order')
     assert [node.key for node in in_order] == [5, 1, 3, 2, 4, 0, 6]
     assert [node.key for node in pre_order] == [2, 1, 5, 3, 0, 4, 6]
 
-    a4 = AVLTree(0,0,backend=backend)
-    for i in range(1,7):
+    a4 = AVLTree()
+    a4.set_tree( ArrayForTrees(TreeNode, 0, backend=backend) )
+    for i in range(0,7):
         a4.tree.append(TreeNode(i, i,backend=backend))
     a4.tree[0].left = 1
     a4.tree[0].right = 2
@@ -250,7 +253,7 @@ def _test_AVLTree(backend):
 
     a5 = AVLTree(is_order_statistic=True,backend=backend)
     if backend==Backend.PYTHON:
-        a5.tree = ArrayForTrees(TreeNode, [
+        a5.set_tree( ArrayForTrees(TreeNode, [
         TreeNode(10, 10),
         TreeNode(5, 5),
         TreeNode(17, 17),
@@ -265,9 +268,9 @@ def _test_AVLTree(backend):
         TreeNode(30, 30),
         TreeNode(13, 13),
         TreeNode(33, 33)
-    ])
+    ]) )
     else:
-        a5.tree = ArrayForTrees(_nodes.TreeNode, [
+        a5.set_tree( ArrayForTrees(_nodes.TreeNode, [
             TreeNode(10, 10,backend=backend),
             TreeNode(5, 5,backend=backend),
             TreeNode(17, 17,backend=backend),
@@ -282,7 +285,7 @@ def _test_AVLTree(backend):
             TreeNode(30, 30,backend=backend),
             TreeNode(13, 13,backend=backend),
             TreeNode(33, 33,backend=backend)
-        ],backend=backend)
+        ],backend=backend) )
 
     a5.tree[0].left, a5.tree[0].right, a5.tree[0].parent, a5.tree[0].height = \
         1, 2, None, 4
@@ -328,30 +331,32 @@ def _test_AVLTree(backend):
     a5.tree[11].size = 2
     a5.tree[12].size = 1
     a5.tree[13].size = 1
-
-    # assert raises(ValueError, lambda: a5.select(0))
-    # assert raises(ValueError, lambda: a5.select(15))
-    # assert a5.rank(-1) is None
-    # def test_select_rank(expected_output):
-    #     output = []
-    #     for i in range(len(expected_output)):
-    #         output.append(a5.select(i + 1).key)
-    #     assert output == expected_output
-
-    #     output = []
-    #     expected_ranks = [i + 1 for i in range(len(expected_output))]
-    #     for i in range(len(expected_output)):
-    #         output.append(a5.rank(expected_output[i]))
-    #     assert output == expected_ranks
-
-    # test_select_rank([2, 3, 5, 9, 10, 11, 12, 13, 15, 17, 18, 20, 30, 33])
+    assert str(a5) == "[(1, 10, 10, 2), (3, 5, 5, 4), (5, 17, 17, 6), (None, 2, 2, 7), (None, 9, 9, None), (8, 12, 12, 9), (10, 20, 20, 11), (None, 3, 3, None), (None, 11, 11, None), (12, 15, 15, None), (None, 18, 18, None), (None, 30, 30, 13), (None, 13, 13, None), (None, 33, 33, None)]"
+
+    assert raises(ValueError, lambda: a5.select(0))
+    assert raises(ValueError, lambda: a5.select(15))
+    
+    assert a5.rank(-1) is None
+    def test_select_rank(expected_output):
+        output = []
+        for i in range(len(expected_output)):
+            output.append(a5.select(i + 1).key)
+        assert output == expected_output
+
+        output = []
+        expected_ranks = [i + 1 for i in range(len(expected_output))]
+        for i in range(len(expected_output)):
+            output.append(a5.rank(expected_output[i]))
+        assert output == expected_ranks
+
+    test_select_rank([2, 3, 5, 9, 10, 11, 12, 13, 15, 17, 18, 20, 30, 33])
     # a5.delete(9)
     # a5.delete(13)
     # a5.delete(20)
 
-    # trav = BinaryTreeTraversal(a5)
-    # in_order = trav.depth_first_search(order='in_order')
-    # pre_order = trav.depth_first_search(order='pre_order')
+    trav = BinaryTreeTraversal(a5)
+    in_order = trav.depth_first_search(order='in_order')
+    pre_order = trav.depth_first_search(order='pre_order')
     # assert [node.key for node in in_order] == [2, 3, 5, 10, 11, 12, 15, 17, 18, 30, 33]
     # assert [node.key for node in pre_order] == [17, 10, 3, 2, 5, 12, 11, 15, 30, 18, 33]
 
@@ -384,7 +389,7 @@ def test_cpp_AVLTree():
     _test_AVLTree(backend=Backend.CPP)
 
 test_AVLTree()
-test_cpp_AVLTree()
+# test_cpp_AVLTree()
 
 def _test_BinaryIndexedTree(backend):