binary indexed tree

fenwick_tree.py

class FenwickTree(object):
  def __init__(self, arr):
    self._one_based_indexed_arr = [0] + arr
    self._bit = [0] * (len(arr)+2)
    self._len = len(self._one_based_indexed_arr)
    for i in range(1, len(self._one_based_indexed_arr)):
      index = i
      val = self._one_based_indexed_arr[i]
      while index < self._len:
        self._bit[index] += val
        index += index&-index

  def update(self, index, val):
    index_plus_one = index+1
    diff = self._one_based_indexed_arr[index_plus_one] - val if self._one_based_indexed_arr[index_plus_one] > val else val - self._one_based_indexed_arr[index_plus_one]
    while index_plus_one < self._len:
      self._bit[index_plus_one] += diff
      index_plus_one += index_plus_one&-index_plus_one

  def get_sum(self, index):
    s = 0
    i_plus_one = index+1
    while i_plus_one:
      s += self._bit[i_plus_one]
      i_plus_one -= i_plus_one&-i_plus_one
    return s


test1 = [2,1,1,3,2,3,4,5,6,7,8,9,10,11,12,13,14,15]
ft = FenwickTree(test1)
prefix_sums = []
s = 0
for x in test1:
  s += x
  prefix_sums.append(s)

assert ft.get_sum(4) == prefix_sums[4]
# add 2 to the first element
ft.update(0, 4)
assert ft.get_sum(4) == prefix_sums[4] + 2
ft.update(7,10)
assert ft.get_sum(7) == prefix_sums[7] + 7