Added Fenwick Trees

pydatastructs/trees/__init__.py

@@ -2,7 +2,8 @@
 
 from . import (
     binary_trees,
-    space_partitioning_trees
+    space_partitioning_trees,
+    fenwick_tree
 )
 
 from .binary_trees import (
@@ -15,6 +16,11 @@
 )
 __all__.extend(space_partitioning_trees.__all__)
 
+from .fenwick_tree import (
+    FenwickTree
+)
+__all__.extend(fenwick_tree.__all__)
+
 from .heaps import (
     BinaryHeap,
     BinomialHeap

pydatastructs/trees/fenwick_tree.py

@@ -0,0 +1,34 @@
+__all__ = ['FenwickTree']
+class FenwickTree():
+    def __init__(self, array):
+        self.size = len(array) + 1
+        self.tree = [0] * (self.size + 1)
+        self.array = array
+        self.flag = [0] * (self.size-1)
+        for index in range(self.size-1):
+            self.update(index, array[index])
+    def update(self, index, value):
+        if(self.flag[index] == 0):
+            self.flag[index] = 1
+            index += 1
+            while(index < self.size):
+                self.tree[index] += value
+                index = index + (index & (-index))
+        else:
+            value = value - self.array[index]
+            index += 1
+            while(index < self.size):
+                self.tree[index] += value
+                index = index + (index & (-index))
+    def get_prefix_sum(self, index):
+        index += 1
+        sum = 0
+        while(index > 0):
+            sum += self.tree[index]
+            index = index - (index & (-index))
+        return sum
+    def get_sum(self, left_index, right_index):
+        if(left_index >= 1):
+            return self.get_prefix_sum(right_index) - self.get_prefix_sum(left_index - 1)
+        else:
+            return self.get_prefix_sum(right_index)

pydatastructs/trees/tests/test_fenwick_tree.py

@@ -0,0 +1,13 @@
+from pydatastructs.trees import FenwickTree
+from pydatastructs.utils.raises_util import raises
+from copy import deepcopy
+def test_FenwickTree():
+    FT = FenwickTree
+    t = FT([1,2,3,4,5,6,7,8,9,10])
+    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(0,9) == 154