Python/data_structures/binary_tree/fenwick_tree.py

29 lines
734 B
Python
Raw Normal View History

2018-10-19 12:48:28 +00:00
class FenwickTree:
2019-10-05 05:14:13 +00:00
def __init__(self, SIZE): # create fenwick tree with size SIZE
2018-10-19 12:48:28 +00:00
self.Size = SIZE
2019-10-05 05:14:13 +00:00
self.ft = [0 for i in range(0, SIZE)]
2018-10-19 12:48:28 +00:00
2019-10-05 05:14:13 +00:00
def update(self, i, val): # update data (adding) in index i in O(lg N)
while i < self.Size:
2018-10-19 12:48:28 +00:00
self.ft[i] += val
i += i & (-i)
2019-10-05 05:14:13 +00:00
def query(self, i): # query cumulative data from index 0 to i in O(lg N)
2018-10-19 12:48:28 +00:00
ret = 0
2019-10-05 05:14:13 +00:00
while i > 0:
2018-10-19 12:48:28 +00:00
ret += self.ft[i]
i -= i & (-i)
return ret
2019-10-05 05:14:13 +00:00
if __name__ == "__main__":
2018-10-19 12:48:28 +00:00
f = FenwickTree(100)
2019-10-05 05:14:13 +00:00
f.update(1, 20)
f.update(4, 4)
print(f.query(1))
print(f.query(3))
print(f.query(4))
2019-10-05 05:14:13 +00:00
f.update(2, -5)
print(f.query(1))
print(f.query(3))