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.

思路：Naive的方法容易想到，复杂度是O(N^2)

一个更好的想法是merge sort，与 Count of Smaller Numbers After Self 类似，http://blog.csdn.net/zjucor/article/details/56667936

关键是利用merge sort过程中 “部分是排序好的” 这一特性，边归并，边累加

Recall count smaller number after self where we encountered the problem

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

Here, after we did the preprocess, we need to solve the problem

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

Therefore the two problems are almost the same. We can use the same technique used in that problem to solve this problem. One solution is merge sort based; another one is Balanced BST based. The time complexity are both O(n log n).

The merge sort based solution counts the answer while doing the merge. During the merge stage, we have already sorted the left half[start, mid) and right half [mid, end). We then iterate through the left half with index i. For each i, we need to find two indices k and j in the right half where

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

Then the number of sums in [lower, upper] is j-k. We also use another index t to copy the elements satisfy sums[t] < sums[i] to a cache in order to complete the merge sort.

Despite the nested loops, the time complexity of the "merge & count" stage is still linear. Because the indices k, j, t will only increase but not decrease, each of them will only traversal the right half once at most. The total time complexity of this divide and conquer solution is then O(n log n).

One other concern is that the sums may overflow integer. So we use long instead.

public class Solution {
	
	int rst = 0, lower, upper, s, t;
	
    public int countRangeSum(int[] nums, int lower, int upper) {
        this.lower = lower;
        this.upper = upper;
        int len = nums.length;
        
        long[] sum = new long[len+1];
        for(int i=1; i<len+1; i++)		sum[i] = sum[i-1] + nums[i-1];
        
        sort(sum, new long[len+1], 0, len);
        
    	return rst;
    }
    
    
    public void sort(long[] a, long[] aux, int lo, int hi) {
    	if(hi <= lo)	return;
    	int mid = lo + (hi - lo) / 2;
    	sort(a, aux, lo, mid);
    	sort(a, aux, mid+1, hi);
    	merge(a, aux, lo, mid, hi);
    }


	private void merge(long[] a, long[] aux, int lo, int mid, int hi) {
		
		for(int k=lo; k<=hi; k++) {
			aux[k] = a[k];
		}
		
		// 记下符合区间的数
		s=mid+1; t=mid+1;			// 放在外面，每次循环在上次的基础上继续进行
		for(int k=lo; k<=mid; k++) {
			while(s<=hi && aux[s]-aux[k]<lower)		s++;
			while(t<=hi && aux[t]-aux[k]<=upper)	t++;
			rst += t - s;
		}
		
		// 归并套路
		int i = lo, j = mid+1;
		for(int k=lo; k<=hi; k++) {
			if(i > mid) {
				a[k] = aux[j++];
			} else if(j > hi) {
				a[k] = aux[i++];
			} else if(aux[j] < aux[i])	{
				a[k] = aux[j++];
			} else	{
				a[k] = aux[i++];
			}
		}
	}
}

虽然有嵌套循环，但是一方面是分治，另外一方面2个while循环的范围越来越小，所以总体来说，感觉会比O(N^2)要好