@@ -24,6 +24,33 @@ static void BinaryIndexedTree_dealloc(BinaryIndexedTree *self) {
24
24
Py_TYPE (self)->tp_free (reinterpret_cast <PyObject*>(self));
25
25
}
26
26
27
+ static PyObject* BinaryIndexedTree_update (BinaryIndexedTree* self, PyObject *args) {
28
+ long index = PyLong_AsLong (PyObject_GetItem (args, PyZero));
29
+ long value = PyLong_AsLong (PyObject_GetItem (args, PyOne));
30
+ long _index = index;
31
+ long _value = value;
32
+ if (PyList_GetItem (self->flag , index) == PyZero) {
33
+ PyList_SetItem (self->flag , index, PyOne);
34
+ index += 1 ;
35
+ while (index < self->array ->_size + 1 ) {
36
+ long curr = PyLong_AsLong (PyList_GetItem (self->tree , index));
37
+ PyList_SetItem (self->tree , index, PyLong_FromLong (curr + value));
38
+ index = index + (index & (-1 *index));
39
+ }
40
+ }
41
+ else {
42
+ value = value - PyLong_AsLong (self->array ->_data [index]);
43
+ index += 1 ;
44
+ while (index < self->array ->_size + 1 ) {
45
+ long curr = PyLong_AsLong (PyList_GetItem (self->tree , index));
46
+ PyList_SetItem (self->tree , index, PyLong_FromLong (curr + value));
47
+ index = index + (index & (-1 *index));
48
+ }
49
+ }
50
+ self->array ->_data [_index] = PyLong_FromLong (_value);
51
+ Py_RETURN_NONE;
52
+ }
53
+
27
54
static PyObject* BinaryIndexedTree___new__ (PyTypeObject* type, PyObject *args, PyObject *kwds) {
28
55
BinaryIndexedTree *self;
29
56
self = reinterpret_cast <BinaryIndexedTree*>(type->tp_alloc (type, 0 ));
@@ -45,13 +72,42 @@ static PyObject* BinaryIndexedTree___new__(PyTypeObject* type, PyObject *args, P
45
72
self->flag = PyList_New (self->array ->_size );
46
73
for (int i=0 ;i<self->array ->_size ;i++){
47
74
PyList_SetItem (self->flag , i, PyZero);
75
+ BinaryIndexedTree_update (self, Py_BuildValue (" (OO)" , PyLong_FromLong (i), self->array ->_data [i]));
48
76
}
49
77
50
78
return reinterpret_cast <PyObject*>(self);
51
79
}
52
80
81
+ static PyObject* BinaryIndexedTree_get_prefix_sum (BinaryIndexedTree* self, PyObject *args) {
82
+ long index = PyLong_AsLong (PyObject_GetItem (args, PyZero));
83
+ index += 1 ;
84
+ long sum = 0 ;
85
+ while (index > 0 ) {
86
+ sum += PyLong_AsLong (PyList_GetItem (self->tree , index));
87
+ index = index - (index & (-1 *index));
88
+ }
89
+
90
+ return PyLong_FromLong (sum);
91
+ }
92
+
93
+ static PyObject* BinaryIndexedTree_get_sum (BinaryIndexedTree* self, PyObject *args) {
94
+ long left_index = PyLong_AsLong (PyObject_GetItem (args, PyZero));
95
+ long right_index = PyLong_AsLong (PyObject_GetItem (args, PyOne));
96
+ if (left_index >= 1 ) {
97
+ long l1 = PyLong_AsLong (BinaryIndexedTree_get_prefix_sum (self, Py_BuildValue (" (O)" , PyLong_FromLong (right_index))));
98
+ long l2 = PyLong_AsLong (BinaryIndexedTree_get_prefix_sum (self, Py_BuildValue (" (O)" , PyLong_FromLong (left_index - 1 ))));
99
+ return PyLong_FromLong (l1 - l2);
100
+ }
101
+ else {
102
+ return BinaryIndexedTree_get_prefix_sum (self, Py_BuildValue (" (O)" , PyLong_FromLong (right_index)));
103
+ }
104
+ }
105
+
53
106
54
107
static struct PyMethodDef BinaryIndexedTree_PyMethodDef[] = {
108
+ {" update" , (PyCFunction) BinaryIndexedTree_update, METH_VARARGS, NULL },
109
+ {" get_prefix_sum" , (PyCFunction) BinaryIndexedTree_get_prefix_sum, METH_VARARGS, NULL },
110
+ {" get_sum" , (PyCFunction) BinaryIndexedTree_get_sum, METH_VARARGS, NULL },
55
111
{NULL }
56
112
};
57
113
0 commit comments