Count of Range Sum [Hard] — Merge Sort on Prefix Sums

Problem Statement

Count the number of range sums sum(i,j) (0 ≤ i ≤ j ≤ n-1) that lie in [lower, upper].

Input: nums=[-2,5,-1], lower=-2, upper=2 → Output: 3

Intuition

Build prefix sums. For each P[i], we need j < i where lower <= P[i]-P[j] <= upper, i.e. P[i]-upper <= P[j] <= P[i]-lower. Use merge sort: during merge, count valid pairs across halves with two pointers.

Solutions

Python

def countRangeSum(nums: list[int], lower: int, upper: int) -> int:
    prefix = [0]
    for n in nums:
        prefix.append(prefix[-1] + n)

    def merge_sort(arr):
        if len(arr) <= 1:
            return arr, 0
        mid = len(arr) // 2
        left, lc = merge_sort(arr[:mid])
        right, rc = merge_sort(arr[mid:])
        count = lc + rc
        # Count: for each right[j], count left[i] in [right[j]-upper, right[j]-lower]
        j = k = 0
        for r in right:
            while j < len(left) and left[j] < r - upper: j += 1
            while k < len(left) and left[k] <= r - lower: k += 1
            count += k - j
        # Merge
        merged = []
        i = p = 0
        while i < len(left) and p < len(right):
            if left[i] <= right[p]: merged.append(left[i]); i += 1
            else: merged.append(right[p]); p += 1
        merged.extend(left[i:]); merged.extend(right[p:])
        return merged, count

    _, total = merge_sort(prefix)
    return total

Complexity

Time: O(n log n)
Space: O(n)