Skip to content

Instantly share code, notes, and snippets.

@ssanin82
Last active December 16, 2017 08:18
Show Gist options
  • Save ssanin82/b031ce580157dd46f9a91dca18a4474f to your computer and use it in GitHub Desktop.
Save ssanin82/b031ce580157dd46f9a91dca18a4474f to your computer and use it in GitHub Desktop.
class SegmentTree:
def __init__(self, n, a, combine):
self.n = n
self.tree = [0] * 2 * self.n
self.op = combine
for i in range(n):
self.tree[n + i] = a[i]
for i in range(self.n - 1, 0, -1):
self.tree[i] = self.op(self.tree[i << 1], self.tree[i << 1 | 1])
def modify(self, p, value):
p += self.n
self.tree[p] = value
while p > 1:
self.tree[p >> 1] = self.op(self.tree[p], self.tree[p ^ 1]) # can be left or right child => ^1 to flip
p >>= 1
def query(self, l, r): # query [l, r)
res, l, r = 0, n + l, n + r
while l < r:
if l % 2: # right child
res = self.op(res, self.tree[l]) # add this value, not parent's
l = (l + 1) // 2 # move to the right of its parent
else: # left child
l >>= 1 # move to parent
if r % 2: # right child
r -= 1 # go to left child
res = self.op(res, self.tree[r]) # add left child
r >>= 1 # move to parent
return res
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment