C++ backend for Binary Indexed Trees completed

Author: Kishan-Ved

pydatastructs/trees/_backend/cpp/BinaryIndexedTree.hpp

@@ -24,6 +24,33 @@ static void BinaryIndexedTree_dealloc(BinaryIndexedTree *self) {
     Py_TYPE(self)->tp_free(reinterpret_cast<PyObject*>(self));
 }
 
+static PyObject* BinaryIndexedTree_update(BinaryIndexedTree* self, PyObject *args) {
+    long index = PyLong_AsLong(PyObject_GetItem(args, PyZero));
+    long value = PyLong_AsLong(PyObject_GetItem(args, PyOne));
+    long _index = index;
+    long _value = value;
+    if (PyList_GetItem(self->flag, index) == PyZero) {
+        PyList_SetItem(self->flag, index, PyOne);
+        index += 1;
+        while (index < self->array->_size + 1) {
+            long curr = PyLong_AsLong(PyList_GetItem(self->tree, index));
+            PyList_SetItem(self->tree, index, PyLong_FromLong(curr + value));
+            index = index + (index & (-1*index));
+        }
+    }
+    else {
+        value = value - PyLong_AsLong(self->array->_data[index]);
+        index += 1;
+        while (index < self->array->_size + 1) {
+            long curr = PyLong_AsLong(PyList_GetItem(self->tree, index));
+            PyList_SetItem(self->tree, index, PyLong_FromLong(curr + value));
+            index = index + (index & (-1*index));
+        }
+    }
+    self->array->_data[_index] = PyLong_FromLong(_value);
+    Py_RETURN_NONE;
+}
+
 static PyObject* BinaryIndexedTree___new__(PyTypeObject* type, PyObject *args, PyObject *kwds) {
     BinaryIndexedTree *self;
     self = reinterpret_cast<BinaryIndexedTree*>(type->tp_alloc(type, 0));
@@ -45,13 +72,42 @@ static PyObject* BinaryIndexedTree___new__(PyTypeObject* type, PyObject *args, P
     self->flag = PyList_New(self->array->_size);
     for(int i=0;i<self->array->_size;i++){
         PyList_SetItem(self->flag, i, PyZero);
+        BinaryIndexedTree_update(self, Py_BuildValue("(OO)", PyLong_FromLong(i), self->array->_data[i]));
     }
 
     return reinterpret_cast<PyObject*>(self);
 }
 
+static PyObject* BinaryIndexedTree_get_prefix_sum(BinaryIndexedTree* self, PyObject *args) {
+    long index = PyLong_AsLong(PyObject_GetItem(args, PyZero));
+    index += 1;
+    long sum = 0;
+    while (index > 0) {
+        sum += PyLong_AsLong(PyList_GetItem(self->tree, index));
+        index = index - (index & (-1*index));
+    }
+
+    return PyLong_FromLong(sum);
+}
+
+static PyObject* BinaryIndexedTree_get_sum(BinaryIndexedTree* self, PyObject *args) {
+    long left_index = PyLong_AsLong(PyObject_GetItem(args, PyZero));
+    long right_index = PyLong_AsLong(PyObject_GetItem(args, PyOne));
+    if (left_index >= 1) {
+        long l1 = PyLong_AsLong(BinaryIndexedTree_get_prefix_sum(self, Py_BuildValue("(O)", PyLong_FromLong(right_index))));
+        long l2 = PyLong_AsLong(BinaryIndexedTree_get_prefix_sum(self, Py_BuildValue("(O)", PyLong_FromLong(left_index - 1))));
+        return PyLong_FromLong(l1 - l2);
+    }
+    else {
+        return BinaryIndexedTree_get_prefix_sum(self, Py_BuildValue("(O)", PyLong_FromLong(right_index)));
+    }
+}
+
 
 static struct PyMethodDef BinaryIndexedTree_PyMethodDef[] = {
+    {"update", (PyCFunction) BinaryIndexedTree_update, METH_VARARGS, NULL},
+    {"get_prefix_sum", (PyCFunction) BinaryIndexedTree_get_prefix_sum, METH_VARARGS, NULL},
+    {"get_sum", (PyCFunction) BinaryIndexedTree_get_sum, METH_VARARGS, NULL},
     {NULL}
 };
 

pydatastructs/trees/tests/test_binary_trees.py

@@ -361,20 +361,25 @@ def test_select_rank(expected_output):
     test_select_rank([])
 
 
-def test_BinaryIndexedTree():
+def _test_BinaryIndexedTree(backend):
 
     FT = BinaryIndexedTree
 
-    t = FT([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], backend=Backend.CPP)
+    t = FT([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], backend=backend)
+
+    assert t.get_sum(0, 2) == 6
+    assert t.get_sum(0, 4) == 15
+    assert t.get_sum(0, 9) == 55
+    t.update(0, 100)
+    assert t.get_sum(0, 2) == 105
+    assert t.get_sum(0, 4) == 114
+    assert t.get_sum(1, 9) == 54
 
-    # assert t.get_sum(0, 2) == 6
-    # assert t.get_sum(0, 4) == 15
-    # assert t.get_sum(0, 9) == 55
-    # t.update(0, 100)
-    # assert t.get_sum(0, 2) == 105
-    # assert t.get_sum(0, 4) == 114
-    # assert t.get_sum(1, 9) == 54
+def test_BinaryIndexedTree():
+    _test_BinaryIndexedTree(Backend.PYTHON)
 
+def test_cpp_BinaryIndexedTree():
+    _test_BinaryIndexedTree(Backend.CPP)
 
 def test_CartesianTree():
     tree = CartesianTree()