327. Count of Range Sum

问题描述

Given an integer array nums, return the number of range sums that lie in [lower, upper] inclusive.Range sum S(i, j) is defined as the sum of the elements in nums between indices i and j (i ≤ j), inclusive.
Note:
A naive algorithm of O(n2) is trivial. You MUST do better than that.
Example:
Given nums = [-2, 5, -1], lower = -2, upper = 2,Return 3.
The three ranges are : [0, 0], [2, 2], [0, 2] and their respective sums are: -2, -1, 2.

问题分析

一道hard题，想了一下，本来想用两个指针移动来做，但由于数组是未排序的，而且还有正有负，我想的方法不可行，因此去网上搜了结题报告。
有两个思路：

1. Merge Sort方法。构建一个sums数组，对它进行merge sort，在merge的过程中抽取解。具体分析可见Xuan's blog，说得非常清楚明白。
   这里简单记录一下思路：在合并sums[left:mid]和sums[mid+1:right]时，两个数组分别已经是有序的，因此可以使用两指针的方法，对于左数组中的每一个元素，在右数组中寻找rl、rr;
   rl：对于左数组中的sums[i]，右数组中第一个不满足sums[rl] - sums[i] < lower的位置；
   rr：对于左数组中的sums[i]，右数组中第一个不满足sums[rr] - sums[i] <= upper的位置；
   那么rr-rl就是右数组元素减sums[i]在[lower，upper]中的个数。
   由于两数组都是递增的，左数组向右移动的过程中，右数组的rl、rr指针也是向右移动的，不会回溯，因此合并部分的复杂度是O(n)。
2. Fenwick Tree方法。之前在做315. Count of Smaller Numbers After Self时也听到过这个名词，但当时没有理解。这次看到一篇讲解很清晰的文章，算是了解了Fenwick Tree究竟是啥。但是如何利用Fenwick Tree解这道题我还没有理解。

AC代码

class Solution(object):
    def countRangeSum(self, nums, lower, upper):
        """
        :type nums: List[int]
        :type lower: int
        :type upper: int
        :rtype: int
        """
        if len(nums) == 0:
            return 0
        sums = [0]
        for i in range(len(nums)):
            sums.append(sums[i] + nums[i])
        return self.count(sums, 0, len(nums), lower, upper)
        
    def count(self, sums, left, right, lower, upper):
        if right - left <= 0:
            return 0
        mid = (left + right) / 2
        rst = self.count(sums, left, mid, lower, upper) + self.count(sums, mid+1, right, lower, upper)
        rl = rr = mid+1
        for i in range(left, mid+1):
            while rl <= right and sums[rl] - sums[i] < lower:
                rl += 1
            while rr <= right and sums[rr] - sums[i] <= upper:
                rr += 1
            rst += rr-rl
        tmp = sums[left:right+1]
        p = 0
        q = mid + 1 -left
        for i in range(left, right+1):
            if p <= mid-left and q <= right-left:
                if tmp[p] < tmp[q]:
                    sums[i] = tmp[p]
                    p += 1
                else:
                    sums[i] = tmp[q]
                    q += 1
            elif p <= mid-left:
                sums[i] = tmp[p]
                p += 1
            else:
                sums[i] = tmp[q]
                q += 1
        return rst

Runtime: 346 ms, which beats 30.51% of python submissions.