Leetcode 327: Count of Range Sum

Given an integer array nums, return the number of range sums that lie in [lower, upper] inclusive.
Range sum S(i, j) is defined as the sum of the elements in nums between indices i and j (i ≤ j), inclusive.

Note:
A naive algorithm of O(n²) is trivial. You MUST do better than that.

Example:

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

Problem Analysis:

We can simply get the O(n²) solution.

class Solution {
public:
  int countRangeSum(vector<int>& nums, int lower, int upper) {
    vector<int> sum(nums.size() + 1, 0);
    for (int i = 0; i < nums.size(); ++i) 
      sum[i + 1] = sum[i] + nums[i];
    int count = 0;
    for (int i = 0; i < sum.size(); ++i) {
      for (int j = i + 1; j < sum.size(); ++j) {
        if (sum[j] - sum[i] >= lower && sum[j] - sum[i] <= upper)
          ++count;
      }
    }
    return count;
  }
};

To optimize time complexity, we try to use some other solution. Firstly we think of using sort, the problem is similar as to sort of sum array. We can use merge sort, in sorting, we count that lower ≤ sum[j]-sum[i]≤upper.

Solution:

class Solution {
public:
  int countRangeSum(vector<int>& nums, int lower, int upper) {
    vector<long long> sum(nums.size() + 1, 0);
    for (int i = 0; i < nums.size(); ++i) 
      sum[i + 1] = sum[i] + nums[i];
    
    return merge(sum, 0, sum.size(), lower, upper);
  }
    
  int merge(vector<long long>& nums, int start, int end, int lower, int upper) {
    if (end - start <= 1) return 0;
        
    int middle = start + (end - start) / 2;
    int count = merge(nums, start, middle, lower, upper) + 
            merge(nums, middle, end, lower, upper);
        
    int small = middle;
    int between = middle;
    for (int i = start; i < middle; ++i) {
      while (small < end && nums[small] - nums[i] < lower) 
        ++small;
      while (between < end && nums[between] - nums[i] <= upper)
        ++between;
      count += between - small;
    }
        
    inplace_merge(nums.begin() + start, nums.begin() + middle, nums.begin() + end);
    return count;
  }
};

Time Complexity: O(nlgn)

Space Complexity: O(n)