327. Count of Range Sum

中午請人吃飯，結果吃多了，好困，有點坐不動了。這題有segment tree其實不是那么好做，答案的解釋一大堆，決定放到明天去理解。今天的任務是做到hard70，大概是可以完成的。

class Solution(object):
    def merge(self, arr1, arr2):
        r, i, j = [], 0, 0
        while i < len(arr1) and j < len(arr2):
            if arr1[i] < arr2[j]:
                r.append(arr1[i])
                i += 1
            else:
                r.append(arr2[j])
                j += 1
        r += arr1[i:] + arr2[j:]
        return r

    def count(self, prefix, start, end, lower, upper):
        if start >= end:
            return 0

        mid = start + (end - start + 1 >> 1)
        count = sum([
            self.count(prefix, s, e, lower, upper)
            for s, e in ((start, mid - 1), (mid, end))
        ])

        l, r = mid, mid
        for i in xrange(start, mid):
            while l <= end and prefix[l] - prefix[i] < lower:
                l += 1
            while r <= end and prefix[r] - prefix[i] <= upper:
                r += 1
            count += r - l

        prefix[start:end + 1] = self.merge(
            prefix[start:mid], prefix[mid:end + 1])

        return count

    def countRangeSum(self, nums, lower, upper):
        n = len(nums)

        prefix = [0] * (n + 1)
        for i, v in enumerate(nums, 1):
            prefix[i] = prefix[i - 1] + v

        return self.count(prefix, 0, n, lower, upper)