hard Count of Range Sum

Problem Statement

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

A range sum S(i, j) is calculated by summing the elements in nums from index i to index j (where i <= j).

Examples

1. Example 1:

   • Input: nums = [2, -1, 2, -3, 4], lower = 0, upper = 2
   • Output: 7
   • Explanation: The valid range sums are S(0, 0) = 0, S(0, 1) = 1, S(2, 2) = 2, S(3, 4) = 1, S(1, 4) = 2, S(0, 2) = 1, and S(0, 3) = 0. All these sums lie within [0, 2].

2. Example 2:

   • Input: nums = [0, 3, -2, 1, -4, 2], lower = -2, upper = 2
   • Output: 15
   • Explanation: There are a total 15 valid ranges whose sum lies within [-2, 2].

3. Example 3:

   • Input: nums = [1, -1, 1, -1, 1],lower = 0, upper = 1
   • Output: 12
   • Explanation: The valid range sums are S(0, 0) = 1, S(2, 2) = 1, S(4, 4) = 1, S(0, 1) = 0, S(1, 2) = 0, S(2, 3) = 0, S(3, 4) = 0, S(0, 2) = 1, S(2, 5) = 1, S(0, 3) = 0, S(1, 4) = 0, S(0, 4) = 1.

Constraints:

• 1 <= nums.length <= 10⁵
• -2³¹ <= nums[i] <= 2³¹ - 1
• -10⁵ <= lower <= upper <= 10⁵
• The answer is guaranteed to fit in a 32-bit integer.

Solution

To solve this problem, we use a Binary Indexed Tree (BIT) to efficiently count the number of valid range sums. The main idea is to keep track of prefix sums, which allows us to calculate the sum of any subarray in constant time by subtracting two prefix sums. As we calculate each prefix sum, we use the BIT to count how many of these prefix sums, when subtracted from the current prefix sum, fall within the [lower, upper] range.

This approach is effective because it reduces the problem of finding subarrays with sums in a certain range to a problem of counting prefix sums within a range. The BIT helps manage the range queries efficiently, which ensures that our solution scales well even with large input sizes. The overall approach is optimal in terms of both time and space complexity compared to a brute-force method.

Step-by-Step Algorithm

1. Initialize Variables:

   • Define n as the length of the input array nums.
   • Initialize ans as 0 to store the count of valid range sums.
   • Create an array pre of size n+1 to store prefix sums. The extra space is used to simplify calculations when considering sums from the start of the array.

2. Calculate Prefix Sums:

   • Loop through the array nums and calculate the prefix sum for each index. Store these prefix sums in the array pre.
   • This helps us easily calculate the sum of any subarray by subtracting two prefix sums.

3. Sort the Prefix Sum Array:

   • Sort the pre array to allow us to use binary search for efficient indexing.
   • Sorting is important because the Binary Indexed Tree (BIT) relies on the positions of these prefix sums for correct counting.

4. Initialize the Binary Indexed Tree (BIT):

   • Create an integer array bit of size pre.length + 2 to act as the BIT.
   • The BIT will help us efficiently count how many prefix sums fall within a certain range, which is key to counting valid subarrays.

5. Iterate Over nums Array:

   • Initialize sum to 0 to track the cumulative sum as we iterate through nums.

   • Update the BIT:

     • Use the bs(sum, pre) operation to find where the current sum fits into the sorted pre array.
     • This position indicates the cumulative sum up to this point. We update the BIT at this index to record that this prefix sum has been encountered.
     • This operation works because the BIT will later be used to count how many such sums are present that fall within the required range.

   • Update the Cumulative Sum:

     • Add the current nums[i] to sum to update the cumulative sum.

   • Count Valid Range Sums:

     • We want to count how many prefix sums lie within [sum - upper, sum - lower].
     • Use bs(sum-lower, pre) to find how many prefix sums are ≤ sum - lower.
     • Use bs(sum-upper-1, pre) to find how many prefix sums are ≤ sum - upper - 1.
     • The difference between these two counts gives us the number of valid subarrays ending at the current position that satisfy the condition lower ≤ sum(i, j) ≤ upper.
     • The operation ans += sum(bit, bs(sum-lower, pre)) - sum(bit, bs(sum-upper-1, pre)); works because it adds the count of valid range sums up to this point.

6. Return the Result:

   • After iterating through the entire array, ans will hold the total number of valid subarrays.
   • Return ans as the final result.

Why Above Steps Work?

• Binary Search (bs(sum, pre)):

  • This function returns the position where the current cumulative sum would fit in the sorted prefix sum array. It effectively maps the cumulative sum to its index in the BIT. By updating the BIT at this position, we record that this prefix sum has occurred, enabling us to count how many times such sums fall within the desired range.

• Counting Valid Sums:

  • By using the BIT to sum up how many prefix sums are within a certain range, we can efficiently count the number of valid subarrays that satisfy the condition lower ≤ sum(i, j) ≤ upper. The difference of sums from the BIT gives us the exact number of valid subarrays ending at the current position.

Algorithm Walkthrough

Input: nums = [2, -1, 2, -3, 4], lower = 0, upper = 2

1. Initialization:

   • n = 5 (length of nums)
   • ans = 0 (to store the count of valid range sums)
   • pre = [0, 0, 0, 0, 0, 0] (prefix sum array initialized with an extra space)

2. Calculate Prefix Sums:

   • The final prefix sum is pre = [0, 2, 1, 3, 0, 4]

3. Sort the Prefix Sum Array:

   • Sort pre resulting in pre = [0, 0, 1, 2, 3, 4]

4. Initialize the Binary Indexed Tree (BIT):

   • bit = [0, 0, 0, 0, 0, 0, 0, 0] (BIT array initialized with extra space)

5. Iterate Over nums Array:

   • Iteration 1 (i = 0):

     • sum = 0 (initial cumulative sum)
     • Update the BIT:
       • bs(sum, pre) finds the position of sum = 0 in pre, which is index 2
       • Update BIT at index 2, so bit = [0, 0, 0, 1, 1, 0, 0, 0]
     • Update Cumulative Sum:
       • sum = sum + nums[0] = 0 + 2 = 2
     • Count Valid Range Sums:
       • bs(sum-lower, pre) = bs(2-0, pre) = 4 (position of sum-lower in pre)
       • bs(sum-upper-1, pre) = bs(2-2-1, pre) = 0 (position of sum-upper-1 in pre)
       • Count = sum(bit, 4) - sum(bit, 0) = 1 - 0 = 1
       • Add 1 to ans, so ans = 1

   • Iteration 2 (i = 1):

     • sum = 2 (from the previous iteration)
     • Update the BIT:
       • bs(sum, pre) finds the position of sum = 2 in pre, which is index 4
       • Update BIT at index 4, so bit = [0, 0, 0, 1, 1, 1, 1, 0]
     • Update Cumulative Sum:
       • sum = sum + nums[1] = 2 + (-1) = 1
     • Count Valid Range Sums:
       • bs(sum-lower, pre) = bs(1-0, pre) = 3 (position of sum-lower in pre)
       • bs(sum-upper-1, pre) = bs(1-2-1, pre) = 0 (position of sum-upper-1 in pre)
       • Count = sum(bit, 3) - sum(bit, 0) = 1 - 0 = 1
       • Add 1 to ans, so ans = 2

   • Iteration 3 (i = 2):

     • sum = 1 (from the previous iteration)
     • Update the BIT:
       • bs(sum, pre) finds the position of sum = 1 in pre, which is index 3
       • Update BIT at index 3, so bit = [0, 0, 0, 1, 2, 1, 1, 0]
     • Update Cumulative Sum:
       • sum = sum + nums[2] = 1 + 2 = 3
     • Count Valid Range Sums:
       • bs(sum-lower, pre) = bs(3-0, pre) = 5 (position of sum-lower in pre)
       • bs(sum-upper-1, pre) = bs(3-2-1, pre) = 2 (position of sum-upper-1 in pre)
       • Count = sum(bit, 5) - sum(bit, 2) = 3 - 1 = 2
       • Add 2 to ans, so ans = 4

   • Iteration 4 (i = 3):

     • sum = 3 (from the previous iteration)
     • Update the BIT:
       • bs(sum, pre) finds the position of sum = 3 in pre, which is index 5
       • Update BIT at index 5, so bit = [0, 0, 0, 1, 2, 1, 2, 0]
     • Update Cumulative Sum:
       • sum = sum + nums[3] = 3 + (-3) = 0
     • Count Valid Range Sums:
       • bs(sum-lower, pre) = bs(0-0, pre) = 2 (position of sum-lower in pre)
       • bs(sum-upper-1, pre) = bs(0-2-1, pre) = 0 (position of sum-upper-1 in pre)
       • Count = sum(bit, 2) - sum(bit, 0) = 1 - 0 = 1
       • Add 1 to ans, so ans = 5

   • Iteration 5 (i = 4):

     • sum = 0 (from the previous iteration)
     • Update the BIT:
       • bs(sum, pre) finds the position of sum = 0 in pre, which is index 2
       • Update BIT at index 2, so bit = [0, 0, 0, 2, 3, 1, 2, 0]
     • Update Cumulative Sum:
       • sum = sum + nums[4] = 0 + 4 = 4
     • Count Valid Range Sums:
       • bs(sum-lower, pre) = bs(4-0, pre) = 6 (position of sum-lower in pre)
       • bs(sum-upper-1, pre) = bs(4-2-1, pre) = 3 (position of sum-upper-1 in pre)
       • Count = sum(bit, 6) - sum(bit, 3) = 5 - 3 = 2
       • Add 2 to ans, so ans = 7

6. Final Result:

   • After all iterations, ans holds the total count of valid range sums, which is 7.
   • Return ans = 7 as the final output.

Code

import java.util.Arrays;

class Solution {

  public int countRangeSum(int[] nums, int lower, int upper) {
    int n = nums.length,
      ans = 0; // Initialize length of nums array and result variable
    long[] pre = new long[n + 1];

    // Calculate the prefix sums for the array
    for (int i = 0; i < n; i++) {
      pre[i + 1] = nums[i] + pre[i];
    }

    Arrays.sort(pre); // Sort the prefix sum array for binary search
    int[] bit = new int[pre.length + 2];
    long sum = 0;

    // Iterate through each element in nums to calculate valid range sums
    for (int i = 0; i < n; i++) {
      update(bit, bs(sum, pre), 1); // Update the BIT with the current cumulative sum
      sum += nums[i];
      // Count the number of valid range sums that lie between lower and upper
      ans +=
        sum(bit, bs(sum - lower, pre)) - sum(bit, bs(sum - upper - 1, pre));
    }
    return ans;
  }

  // Binary search function to find the index of the first number greater than the given sum
  private int bs(long sum, long[] pre) {
    int lo = 0,
      hi = pre.length; // Set initial low and high pointers for binary search
    while (lo < hi) {
      int mid = (lo + hi) >> 1; // Calculate mid index
      if (pre[mid] > sum) {
        hi = mid; // Move high pointer to mid
      } else {
        lo = mid + 1; // Move low pointer to mid + 1
      }
    }
    return lo;
  }

  // Function to update the BIT
  private void update(int[] bit, int idx, int inc) {
    for (++idx; idx < bit.length; idx += idx & -idx) {
      // Update all relevant indices in the BIT
      bit[idx] += inc; // Increment the BIT value at the current index
    }
  }

  private int sum(int[] bit, int idx) {
    int ans = 0;
    for (++idx; idx > 0; idx -= idx & -idx) {
      // Iterate through the BIT to sum up values
      ans += bit[idx]; // Add the BIT value at the current index to the result
    }
    return ans;
  }

  public static void main(String[] args) {
    Solution solution = new Solution();

    // Example 1
    int[] nums1 = { 2, -1, 2, -3, 4 };
    int lower1 = 0,
      upper1 = 2;
    System.out.println(
      "Example 1 Output: " + solution.countRangeSum(nums1, lower1, upper1)
    );

    // Example 2
    int[] nums2 = { 0, 3, -2, 1, -4, 2 };
    int lower2 = -2,
      upper2 = 2;
    System.out.println(
      "Example 2 Output: " + solution.countRangeSum(nums2, lower2, upper2)
    );

    // Example 3
    int[] nums3 = { 1, -1, 1, -1, 1 };
    int lower3 = 0,
      upper3 = 1;
    System.out.println(
      "Example 3 Output: " + solution.countRangeSum(nums3, lower3, upper3)
    );
  }
}

Complexity Analysis

Time Complexity

• Prefix Sum Calculation: The prefix sum is calculated in time, where n is the length of the input array nums.
• Sorting: Sorting the prefix sum array takes time.
• Binary Indexed Tree Operations:
  • The update and sumBIT functions operate in time.
  • There are n elements, and for each element, we perform constant-time operations with the BIT, resulting in time.

Overall, the time complexity is due to the sorting and BIT operations.

Space Complexity

• Prefix Sum Array: The prefix sum array requires space.
• Binary Indexed Tree: The BIT also requires space.
• Auxiliary Space: Additional space is used for storing the variables, but they are of constant space, so the space complexity remains .

Thus, the overall space complexity is .