Segment Tree Query

Question

sum

Thoughts

Solution

def query(self, root, start, end):
        if not root:
            return 0
        rStart = root.start
        rEnd = root.end

        if (rStart > end) or ( rEnd < start):
            return 0

        #should be >= and <=
        #if query range if larger than the segment tree range, just return the count
        if rStart >= start and rEnd <= end:
            return root.count

        leftCount = 0
        rightCount = 0

        mid = (rStart + rEnd) / 2
        if start <= mid:
            if end > mid:
                #some range are in right part
                leftCount = self.query(root.left, start, mid)
            else:
                leftCount = self.query(root.left, start, end)
        if mid < end:
            if start <= mid:
                #some range are in left part
                #right part should start with mid + 1
                rightCount = self.query(root.right, mid + 1, end)
            else:
                rightCount = self.query(root.right, start, end)
        return leftCount + rightCount