Implementation
class SegmentTreeNode(object):
def __init__(self, start=0, end=0, val=0):
self.val = val
self.start = start
self.end = end
self.left = None
self.right = None
class SegmentTree(object):
def __init__(self, start, end, nums):
self.root = self.build(start, end, nums)
def build(self, start, end, nums):
if start > end:
return None
root = SegmentTreeNode(start, end, 0)
if start < end:
mid = start + (end-start)/2
root.left = self.build(start, mid, nums)
root.right = self.build(mid+1, end, nums)
root.val = root.left.val + root.right.val
else:
root.val = nums[start]
return root
def add(self, index, value):
self.add_helper(self.root, index, value)
def add_helper(self, root, index, value):
if root.start == index and root.end == index:
root.val += value
return
mid = (root.start + root.end) / 2
if root.start <= index and index <= mid:
self.add_helper(root.left, index, value)
elif root.end >= index and index > mid:
self.add_helper(root.right, index, value)
root.val = root.left.val + root.right.val
def query(self, start, end):
return self.query_helper(self.root, start, end)
def query_helper(self, root, start, end):
if not root:
return 0
rStart = root.start
rEnd = root.end
if rStart > end or rEnd < start:
return 0
if rStart >= start and rEnd <= end:
return root.val
leftVal = 0
rightVal = 0
mid = (rStart + rEnd) / 2
if start <= mid:
if end > mid:
leftVal = self.query_helper(root.left, start, mid)
else:
leftVal = self.query_helper(root.left, start, end)
if mid < end:
if start <= mid:
rightVal = self.query_helper(root.right, mid+1, end)
else:
rightVal = self.query_helper(root.right, start, end)
return leftVal + rightVal