327. Count of Range Sum - jiejackyzhang/leetcode-note GitHub Wiki

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.

解题思路为：

naive solution

用sums数组保存0~i的sum，然后检查每一个range的sum，时间复杂度为O(n^2)。

Merge Sort

在"count smaller number after self"问题中，我们解决了这样的问题：

count[i] = count of nums[j] - nums[i] < 0 with j > i

这里，我们要解决的问题可转化为

count[i] = count of a <= S[j] - S[i] <= b with j > i
ans = sum(count[:])

因此这两个问题几乎相同。

我们可以采用merge sort based solution来解。 在做merge时，我们已经有了两个sorted list：[start, mid)和[mid, end)。 我们要做的就是判断起点在left half，终点在right half的range有多少个。 顺序扫描left half，对下标i，我们要在right half中找到两个下标k和j：

  • j is the first index satisfy sums[j] - sums[i] > upper and
  • k is the first index satisfy sums[k] - sums[i] >= lower.

则range sum在[lower, upper]中的有j-k个。 同时用下标t来copy所有sums[t] < sums[i]的元素到cache数组，来完成merge sort。 merge sort一共有logn次iteration，每个iteration的时间是n，因此总的时间复杂度为O(nlogn)。

public class Solution {
    public int countRangeSum(int[] nums, int lower, int upper) {
        int n = nums.length;
        long[] sums = new long[n+1];
        for(int i = 0; i < n; i++) {
            sums[i+1] = sums[i] + nums[i];
        }
        return countWhileMergeSort(sums, 0, n+1, lower, upper);
    }
    
    private int countWhileMergeSort(long[] sums, int start, int end, int lower, int upper) {
        if(end - start <= 1) return 0;
        int mid = start + (end - start) / 2;
        int count = countWhileMergeSort(sums, start, mid, lower, upper) 
                  + countWhileMergeSort(sums, mid, end, lower, upper);
        int j = mid, k = mid, t = mid;
        long[] cache = new long[end-start];
        for(int i = start, r = 0; i < mid; i++, r++) {
            while(k < end && sums[k] - sums[i] < lower) k++;
            while(j < end && sums[j] - sums[i] <= upper) j++;
            while(t < end && sums[t] < sums[i]) cache[r++] = sums[t++];
            cache[r] = sums[i];
            count += j - k;
        }
        System.arraycopy(cache, 0, sums, start, t-start);
        return count;
    }
}
⚠️ **GitHub.com Fallback** ⚠️