Skip to content

Instantly share code, notes, and snippets.

@m00nlight
Last active June 30, 2019 06:10
Show Gist options
  • Save m00nlight/cf89e14d93ed69c204f8 to your computer and use it in GitHub Desktop.
Save m00nlight/cf89e14d93ed69c204f8 to your computer and use it in GitHub Desktop.
python binary index tree
from __future__ import division
class BinaryIndexTree:
def __init__(self, n):
self.sz = n
self.vals = [0] * (n + 1)
def update(self, idx, delta):
"add c to the value at index idx"
while idx <= self.sz:
self.vals[idx] += delta
idx += idx & (-idx)
def accumulate(self, idx):
"get sum from the start to the index of idx"
ret = 0
while idx > 0:
ret += self.vals[idx]
idx -= idx & (-idx)
return ret
def range_sum(self, start, end):
"Calculate a[start], a[start+1], ... a[end]"
return self.accumulate(end) - self.accumulate(start - 1)
def test():
bit = BinaryIndexTree(10)
assert bit.range_sum(1, 5) == 0
bit.update(1, 3)
bit.update(4, 6)
assert bit.range_sum(1, 4) == 9
bit.update(3, 3)
assert bit.range_sum(1, 4) == 12
assert bit.range_sum(3, 4) == 9
bit.update(3, -2)
assert bit.range_sum(3, 3) == 1
assert bit.range_sum(1, 4) == 10
return 'test pass'
if __name__ == '__main__':
print test()
@m00nlight
Copy link
Author

You maked a mistake. line 11 should be while idx<=self.sz, your index start from 0

@Haozun Thanks for pointing out the problem. Fix the mistake. But I think that's because my index start from 1 not 0.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment