原題

對于一個數(shù)組，我們可以對其建立一棵線段樹, 每個結點存儲一個額外的值count來代表這個結點所指代的數(shù)組區(qū)間內(nèi)的元素個數(shù). (數(shù)組中并不一定每個位置上都有元素)
實現(xiàn)一個query的方法，該方法接受三個參數(shù)root, start和end, 分別代表線段樹的根節(jié)點和需要查詢的區(qū)間，找到數(shù)組中在區(qū)間[start, end]內(nèi)的元素個數(shù)。

對于數(shù)組 [0, 2, 3], 對應的線段樹為：

                     [0, 3, count=3]
                     /             \
          [0,1,count=1]             [2,3,count=2]
          /         \               /            \
   [0,0,count=1] [1,1,count=0] [2,2,count=1], [3,3,count=1]

query(1, 1), return 0
query(1, 2), return 1
query(2, 3), return 2
query(0, 2), return 2

解題思路

• 本題為值型線段樹的Query
• 如果查詢區(qū)間 == 節(jié)點表示區(qū)間 => 直接返回root.count
• 如果查詢區(qū)間被節(jié)點表示的左/右區(qū)間包含 => 遞歸搜索左/右區(qū)間
• 如果查詢區(qū)間和結點表示的區(qū)間相交不相等 => 分裂遞歸搜索左/右區(qū)間
• 典型的divide & conquer
• 最后返回 leftCount + rightCount

完整代碼

"""
Definition of SegmentTreeNode:
class SegmentTreeNode:
    def __init__(self, start, end, count):
        self.start, self.end, self.count = start, end, count
        self.left, self.right = None, None
"""

class Solution: 
    # @param root, start, end: The root of segment tree and 
    #                          an segment / interval
    # @return: The count number in the interval [start, end] 
    def query(self, root, start, end):
        if start > end or root == None:
            return 0
            
        if root.start >= start and root.end <= end:
            return root.count
            
        mid = root.start + (root.end - root.start) / 2
        leftCount, rightCount = 0, 0
        if start <= mid:
            if mid < end:
                leftCount = self.query(root.left, start, mid)
            else:
                leftCount = self.query(root.left, start, end)
        if mid < end:
            if start <= mid:
                rightCount = self.query(root.right, mid + 1, end)
            else:
                rightCount = self.query(root.right, start, end)
        return leftCount + rightCount