Count of Range Sum — Merge Sort on Prefix Sums

# Python — merge sort
def countRangeSum(nums, lower, upper):
    prefix = [0]
    for x in nums: prefix.append(prefix[-1]+x)

    def merge_count(arr):
        if len(arr) <= 1: return 0
        mid = len(arr)//2
        count = merge_count(arr[:mid]) + merge_count(arr[mid:])
        left, right = arr[:mid], arr[mid:]
        j = k = 0
        for l in left:
            while j < len(right) and right[j] - l < lower: j += 1
            while k < len(right) and right[k] - l <= upper: k += 1
            count += k - j
        # Merge sorted
        arr[:] = sorted(arr)  # or use standard merge
        return count

    return merge_count(prefix)

// Java
int lower, upper;
public int countRangeSum(int[] nums, int lower, int upper) {
    this.lower=lower; this.upper=upper;
    long[] prefix=new long[nums.length+1];
    for (int i=0;i<nums.length;i++) prefix[i+1]=prefix[i]+nums[i];
    return mergeCount(prefix,0,prefix.length);
}
int mergeCount(long[] p, int lo, int hi) {
    if (hi-lo<=1) return 0;
    int mid=(lo+hi)/2;
    int cnt=mergeCount(p,lo,mid)+mergeCount(p,mid,hi);
    int j=mid, k=mid;
    for (int i=lo;i<mid;i++) {
        while(j<hi&&p[j]-p[i]<lower) j++;
        while(k<hi&&p[k]-p[i]<=upper) k++;
        cnt+=k-j;
    }
    Arrays.sort(p,lo,hi);
    return cnt;
}