327. Count of Range Sum – Leetcode Solutions

Example 1:

Input: nums = [-2,5,-1], lower = -2, upper = 2
Output: 3
Explanation: The three ranges are: [0,0], [2,2], and [0,2] and their respective sums are: -2, -1, 2.

Example 2:

Input: nums = [0], lower = 0, upper = 0
Output: 1

class Solution:
    def countRangeSum(self, nums: List[int], lower: int, upper: int) -> int:
        # Step 1: Compute prefix sums
        # prefix_sum[i] is the sum of nums[0] to nums[i-1]
        prefix_sums = [0] * (len(nums) + 1)
        for i in range(1, len(nums) + 1):
            prefix_sums[i] = prefix_sums[i - 1] + nums[i - 1]

        # Step 2: Define a helper function for merge sort and count the valid range sums
        def merge_sort_and_count(left, right):
            # Base case: only one element
            if left >= right:
                return 0
            
            mid = (left + right) // 2
            count = merge_sort_and_count(left, mid) + merge_sort_and_count(mid + 1, right)

            # Step 3: Count valid range sums in the two sorted halves
            j = k = t = mid + 1
            temp = []
            for i in range(left, mid + 1):
                # Find the range [prefix_sums[i] - upper, prefix_sums[i] - lower] in the right half
                while j <= right and prefix_sums[j] - prefix_sums[i] < lower:
                    j += 1
                while k <= right and prefix_sums[k] - prefix_sums[i] <= upper:
                    k += 1
                count += k - j

                # Step 4: Merge the two halves
                # Adding elements from the right half in sorted order
                while t <= right and prefix_sums[t] < prefix_sums[i]:
                    temp.append(prefix_sums[t])
                    t += 1
                temp.append(prefix_sums[i])

            # Copy the sorted elements back to the prefix_sums array
            prefix_sums[left:left + len(temp)] = temp

            return count

        # Step 5: Start the merge sort process and return the result
        return merge_sort_and_count(0, len(prefix_sums) - 1)

#include <vector>
#include <algorithm>
using namespace std;

class Solution {
public:
    // Function to count the number of range sums that lie in [lower, upper]
    int countRangeSum(vector<int>& nums, int lower, int upper) {
        // Prefix sums array to store cumulative sum up to each index
        vector<long long> prefixSums(nums.size() + 1, 0);
        
        // Create prefix sum for each index
        for (int i = 0; i < nums.size(); ++i) {
            prefixSums[i + 1] = prefixSums[i] + nums[i];
        }
        
        // Use a helper function to count range sums using divide and conquer (merge sort)
        return countAndMergeSort(prefixSums, 0, prefixSums.size(), lower, upper);
    }

private:
    // This helper function uses divide and conquer (merge sort) to count the valid range sums
    int countAndMergeSort(vector<long long>& sums, int start, int end, int lower, int upper) {
        // If the range has only one element, return 0 as there is no range to form
        if (end - start <= 1) return 0;
        
        // Calculate the mid-point to divide the array
        int mid = (start + end) / 2;
        
        // Recursively count range sums in left and right halves
        int count = countAndMergeSort(sums, start, mid, lower, upper) + 
                    countAndMergeSort(sums, mid, end, lower, upper);
        
        // Now count valid range sums that span both halves
        int j = mid, k = mid, t = mid;
        vector<long long> temp;
        
        for (int i = start; i < mid; ++i) {
            // Find the bounds [lower, upper] for range sums
            while (k < end && sums[k] - sums[i] < lower) ++k;
            while (j < end && sums[j] - sums[i] <= upper) ++j;
            
            // Count how many sums are within [lower, upper]
            count += j - k;
            
            // Merge the two halves while maintaining sorted order
            while (t < end && sums[t] < sums[i]) temp.push_back(sums[t++]);
            temp.push_back(sums[i]);
        }
        
        // Add the remaining elements from the right half (if any)
        while (t < end) temp.push_back(sums[t++]);
        
        // Copy back the merged result to the original array
        copy(temp.begin(), temp.end(), sums.begin() + start);
        
        return count;
    }
};

var countRangeSum = function(nums, lower, upper) {
    // Helper function to perform the modified merge sort
    function mergeSort(sum, start, end) {
        // Base case: if the segment is one element, return 0 (no range sums here)
        if (end - start <= 1) return 0;

        // Divide step: split the array segment into two halves
        const mid = Math.floor((start + end) / 2);
        let count = mergeSort(sum, start, mid) + mergeSort(sum, mid, end);

        // Count valid range sums within the current segment
        let j = mid, k = mid, t = mid;
        const cache = [];

        for (let i = start, r = 0; i < mid; i++) {
            // Find the correct position where the range sum becomes valid within [lower, upper]
            while (k < end && sum[k] - sum[i] < lower) k++;
            while (j < end && sum[j] - sum[i] <= upper) j++;

            // Merge the two halves back together
            while (t < end && sum[t] < sum[i]) cache[r++] = sum[t++];
            cache[r++] = sum[i];

            // Count the number of valid sums in the current range
            count += j - k;
        }

        // Merge back the sorted part
        for (let i = 0; i < t - start; i++) {
            sum[start + i] = cache[i];
        }

        return count;
    }

    // Prefix sums array (sum[i] represents the sum of nums from index 0 to i-1)
    const sum = [0];
    for (let i = 0; i < nums.length; i++) {
        sum.push(sum[sum.length - 1] + nums[i]);
    }

    // Call the mergeSort function on the whole prefix sum array
    return mergeSort(sum, 0, sum.length);
};

class Solution {
    
    // This function counts the number of range sums that fall between 'lower' and 'upper'
    public int countRangeSum(int[] nums, int lower, int upper) {
        
        // We first calculate the prefix sums, which will help us quickly compute the sum for any subarray.
        // Prefix sum is used to store the cumulative sum up to each index.
        long[] prefixSums = new long[nums.length + 1];
        
        // Fill the prefixSums array
        for (int i = 0; i < nums.length; i++) {
            prefixSums[i + 1] = prefixSums[i] + nums[i];
        }

        // Call the helper function that uses merge sort to efficiently count the range sums
        return countAndMergeSort(prefixSums, 0, prefixSums.length - 1, lower, upper);
    }

    // This function performs merge sort and counts the valid range sums
    private int countAndMergeSort(long[] prefixSums, int left, int right, int lower, int upper) {
        
        // Base case: when there's only one element, no range sum can be formed
        if (left >= right) return 0;

        // Divide the array into two halves
        int mid = left + (right - left) / 2;

        // Recursively count the range sums in both halves and merge the results
        int count = countAndMergeSort(prefixSums, left, mid, lower, upper)
                    + countAndMergeSort(prefixSums, mid + 1, right, lower, upper);

        // Now count the valid range sums across the two halves
        int rangeCount = 0;
        int j = mid + 1, k = mid + 1;
        for (int i = left; i <= mid; i++) {
            // Find the first 'k' such that prefixSums[k] - prefixSums[i] >= lower
            while (k <= right && prefixSums[k] - prefixSums[i] < lower) k++;

            // Find the first 'j' such that prefixSums[j] - prefixSums[i] > upper
            while (j <= right && prefixSums[j] - prefixSums[i] <= upper) j++;

            // The number of valid sums for this 'i' is (j - k)
            rangeCount += j - k;
        }

        // Merge the two halves (standard merge sort step)
        merge(prefixSums, left, mid, right);

        // Return the total count of range sums
        return count + rangeCount;
    }

    // Merge function to merge two sorted halves
    private void merge(long[] prefixSums, int left, int mid, int right) {
        // Temporary array to store the merged result
        long[] temp = new long[right - left + 1];

        int i = left, j = mid + 1, t = 0;

        // Standard merging of two sorted arrays
        while (i <= mid && j <= right) {
            if (prefixSums[i] <= prefixSums[j]) {
                temp[t++] = prefixSums[i++];
            } else {
                temp[t++] = prefixSums[j++];
            }
        }

        // Copy remaining elements from the left half
        while (i <= mid) {
            temp[t++] = prefixSums[i++];
        }

        // Copy remaining elements from the right half
        while (j <= right) {
            temp[t++] = prefixSums[j++];
        }

        // Copy the sorted array back into the original prefixSums array
        System.arraycopy(temp, 0, prefixSums, left, temp.length);
    }
}