Problem Description

Given an integer array nums and two integers lower and upper, return the number of range sums that lie in [lower, upper] inclusive.

A range sum S(i, j) is defined as the sum of the elements in nums between indices i and j (i.e., nums[i] + nums[i + 1] + ... + nums[j]) where 0 <= i <= j < nums.length.

• Each valid pair (i, j) must satisfy lower <= S(i, j) <= upper.
• The solution should efficiently count all such valid pairs.
• All elements of nums can be negative, zero, or positive.

Thought Process

The problem asks us to count the number of subarrays whose sums fall within a given range. The simplest way is to try every possible subarray, compute its sum, and check if it falls in the range. This brute-force approach checks all O(n^2) subarrays, which is too slow for large arrays.

To optimize, we notice that subarray sums can be computed efficiently using prefix sums. The sum from i to j is prefix[j + 1] - prefix[i]. The challenge now is to count, for each prefix[j+1], how many previous prefix[i] values satisfy lower <= prefix[j+1] - prefix[i] <= upper.

This leads us to techniques like modified merge sort or using Binary Indexed Trees or balanced BSTs for efficient counting.

Solution Approach

• Step 1: Compute Prefix Sums
  • Create an array prefix where prefix[0] = 0 and prefix[i+1] = prefix[i] + nums[i].
  • This allows us to compute any subarray sum in constant time.
• Step 2: Use Modified Merge Sort to Count Valid Ranges
  • We recursively sort the prefix array, and during the merge step, count the number of pairs (i, j) such that lower <= prefix[j] - prefix[i] <= upper.
  • For each left prefix, use two pointers on the right half to find the range of valid right prefixes.
  • This counting is done before merging to preserve the sorted order needed for efficient counting.
• Step 3: Return the Total Count
  • The total count is accumulated during the merge steps.
• Why Merge Sort?
  • Sorting enables us to efficiently count ranges using the properties of sorted arrays, reducing the time complexity to O(n log n).

Example Walkthrough

Suppose nums = [-2, 5, -1], lower = -2, upper = 2.

1. Compute Prefix Sums:
   • prefix = [0, -2, 3, 2]
2. Find All Valid Range Sums:
   • Check all pairs (i, j) where i < j:
   • prefix[1] - prefix[0] = -2 - 0 = -2 (valid)
   • prefix[2] - prefix[0] = 3 - 0 = 3 (invalid)
   • prefix[2] - prefix[1] = 3 - (-2) = 5 (invalid)
   • prefix[3] - prefix[0] = 2 - 0 = 2 (valid)
   • prefix[3] - prefix[1] = 2 - (-2) = 4 (invalid)
   • prefix[3] - prefix[2] = 2 - 3 = -1 (valid)
   • Total valid range sums: 3
3. How Merge Sort Counts:
   • During merge steps, for each left prefix, we use two pointers to count right prefixes that form a valid range sum.

Code Implementation

class Solution:
    def countRangeSum(self, nums, lower, upper):
        def sort(lo, hi):
            if hi - lo <= 1:
                return 0
            mid = (lo + hi) // 2
            count = sort(lo, mid) + sort(mid, hi)
            j = k = mid
            for left in prefix[lo:mid]:
                while k < hi and prefix[k] - left < lower:
                    k += 1
                while j < hi and prefix[j] - left <= upper:
                    j += 1
                count += j - k
            prefix[lo:hi] = sorted(prefix[lo:hi])
            return count

        prefix = [0]
        for num in nums:
            prefix.append(prefix[-1] + num)
        return sort(0, len(prefix))
    

class Solution {
public:
    int countRangeSum(vector<int>& nums, int lower, int upper) {
        vector<long> prefix(nums.size() + 1, 0);
        for (int i = 0; i < nums.size(); ++i)
            prefix[i + 1] = prefix[i] + nums[i];
        return sortCount(prefix, 0, prefix.size(), lower, upper);
    }

    int sortCount(vector<long>& prefix, int lo, int hi, int lower, int upper) {
        if (hi - lo <= 1) return 0;
        int mid = (lo + hi) / 2;
        int count = sortCount(prefix, lo, mid, lower, upper) + sortCount(prefix, mid, hi, lower, upper);
        int j = mid, k = mid;
        for (int i = lo; i < mid; ++i) {
            while (k < hi && prefix[k] - prefix[i] < lower) ++k;
            while (j < hi && prefix[j] - prefix[i] <= upper) ++j;
            count += j - k;
        }
        inplace_merge(prefix.begin() + lo, prefix.begin() + mid, prefix.begin() + hi);
        return count;
    }
};
    

class Solution {
    public int countRangeSum(int[] nums, int lower, int upper) {
        long[] prefix = new long[nums.length + 1];
        for (int i = 0; i < nums.length; ++i)
            prefix[i + 1] = prefix[i] + nums[i];
        return sortCount(prefix, 0, prefix.length, lower, upper);
    }

    private int sortCount(long[] prefix, int lo, int hi, int lower, int upper) {
        if (hi - lo <= 1) return 0;
        int mid = (lo + hi) / 2;
        int count = sortCount(prefix, lo, mid, lower, upper) + sortCount(prefix, mid, hi, lower, upper);
        int j = mid, k = mid;
        for (int i = lo; i < mid; ++i) {
            while (k < hi && prefix[k] - prefix[i] < lower) ++k;
            while (j < hi && prefix[j] - prefix[i] <= upper) ++j;
            count += j - k;
        }
        Arrays.sort(prefix, lo, hi);
        return count;
    }
}
    

var countRangeSum = function(nums, lower, upper) {
    let prefix = [0];
    for (let num of nums) {
        prefix.push(prefix[prefix.length - 1] + num);
    }
    function sort(lo, hi) {
        if (hi - lo <= 1) return 0;
        let mid = Math.floor((lo + hi) / 2);
        let count = sort(lo, mid) + sort(mid, hi);
        let j = mid, k = mid;
        for (let i = lo; i < mid; ++i) {
            while (k < hi && prefix[k] - prefix[i] < lower) k++;
            while (j < hi && prefix[j] - prefix[i] <= upper) j++;
            count += j - k;
        }
        let sorted = [];
        let l = lo, r = mid;
        while (l < mid && r < hi) {
            if (prefix[l] < prefix[r]) sorted.push(prefix[l++]);
            else sorted.push(prefix[r++]);
        }
        while (l < mid) sorted.push(prefix[l++]);
        while (r < hi) sorted.push(prefix[r++]);
        for (let i = 0; i < sorted.length; ++i) {
            prefix[lo + i] = sorted[i];
        }
        return count;
    }
    return sort(0, prefix.length);
};
    

Time and Space Complexity

• Brute-force Approach:
  • Time Complexity: O(n^2) because it checks every possible subarray.
  • Space Complexity: O(1) extra, or O(n) if using prefix sums for speedup.
• Optimized (Modified Merge Sort) Approach:
  • Time Complexity: O(n \log n) because each divide-and-conquer step processes O(n) elements over O(\log n) levels.
  • Space Complexity: O(n) for the prefix sum and temporary arrays used in merge sort.

Summary

The problem of counting range sums is efficiently solved by combining prefix sums with a modified merge sort. Prefix sums allow quick calculation of subarray sums, and merge sort enables us to count valid ranges in O(n \log n) time by leveraging the sorted order. This approach is both elegant and practical, turning a seemingly quadratic problem into an efficient divide-and-conquer solution.