Count of Range Sum - Teeeeeebag/LeetCode GitHub Wiki

1. the Key is cnt+= t-j

public class Solution {
    private int cnt = 0;
    private void mergeSort(long[] sums, int s, int e, long lower, long upper){
        //System.out.println(s + " s e " + e);
        if (s >= e){ 
            return;
        }   
        int mid = s + (e-s)/2;
        mergeSort(sums, s, mid, lower, upper);
        mergeSort(sums, mid+1, e, lower, upper);
        long[] tmp = new long[e-s+1];
        //merge
        int k = mid+1, p = 0, j = mid + 1, t = mid + 1;
        for (int i=s; i<=mid; ++i){
            //The key here is that you don't have to start from mid+1 for each i in [s, mid]
            //sums[j] - sums[i] when sums[i] increase, you just need to try pushing j further    
            while (j<=e && sums[j] - sums[i] < lower){
                ++j;
            }   
            while(t<=e && sums[t] - sums[i] <= upper){
                //System.out.println(sums[j] + " " + sums[i]);
                ++t;
            }   
            while (k <=e && sums[k] < sums[i]){
                tmp[p++] = sums[k++];
            }   
            cnt += t-j; // t-j is the key here, you don't have to reset t and j to mid+1 for each iteration.
                        // This is the key of this algorithm. If you reset t and j to mid+1, time complexity would be n^2logn   
            tmp[p++] = sums[i];
        }   
        System.arraycopy(tmp, 0, sums, s, k-s); // exactly k-s elements to copy, k may be less than e!!!
    }   
    public int countRangeSum(int[] nums, int lower, int upper) {
        if (nums.length == 0){ 
            return 0;
        }   
        long[] sums = new long[nums.length];
        for (int i=0; i<nums.length; ++i){
            if (i == 0){ 
                sums[i] = nums[i];
            } else {
                sums[i] = sums[i-1] + nums[i];
            }   
            if (sums[i] >= lower && sums[i] <= upper){
                ++cnt;
            }
        }   
        //System.out.println(Arrays.toString(sums) + " " + cnt);
        mergeSort(sums, 0, sums.length-1, (long) lower, (long) upper);
        return cnt;
    }   
}