Question
Formatted question description: https://leetcode.ca/all/327.html
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^2) 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.
@taginterval
Algorithm
The most intuitive idea of this question is to enumerate all the subarrays and then check whether they are within the required value range. The time complexity of this method is O(n^2), and it will obviously time out.
What is the easiest thing to think of when seeing this kind of topic? Two Pointers! Yes, but only using Two Pointers is definitely not enough for this problem. Based on the ideas of Two Pointers, we can find a better solution by merging and sorting.
Here our sorting object is the prefix sum array. In the merge stage of merge sort, we have left and right arrays, and both the left and right arrays are sorted, so we can traverse the left array with i, j, k The two pointers are searched in the right array, so that:
sums[j]sums[i] <upper
sums[k]sums[i] >= lower
So at this time, we have found jk subarrays that meet the requirements.
Since the left and right arrays are all sorted, when i is incremented, the positions of j and k need not be scanned from the beginning.
One last thing to note is that in order to prevent overflow, our vector contains long long elements.
Code
Java

import java.util.Collection; import java.util.TreeMap; public class Count_of_Range_Sum { public static void main(String[] args) { Count_of_Range_Sum out = new Count_of_Range_Sum(); Solution s = out.new Solution(); System.out.println(s.countRangeSum(new int[]{2,5,1}, 2, 2)); } public class Solution { public int countRangeSum(int[] nums, int lower, int upper) { if (nums == null  nums.length == 0) { return 0; } long[] sums = new long[nums.length]; sums[0] = nums[0]; for(int i = 1; i < nums.length; i++) { sums[i] = sums[i1] + nums[i]; } int total = 0; TreeMap<Long, Integer> treemap = new TreeMap<>(); for(int i = 0; i < nums.length; i++) { if (lower <= sums[i] && sums[i] <= upper) { total ++; } Collection<Integer> search = treemap .subMap(sums[i]  upper, true, sums[i]lower, true) .values(); for(Integer count: search) { total += count; } // update treemap Integer count = treemap.get(sums[i]); if (count == null) { count = 1; } else { count ++; } treemap.put(sums[i], count); } return total; } } }

Todo

class Solution(object): def countRangeSum(self, nums, lower, upper): """ :type nums: List[int] :type lower: int :type upper: int :rtype: int """ def update(b, i, delta): while i < len(b): b[i] += delta i += (i & i) def sumRange(b, i): ret = 0 while i > 0: ret += b[i] i = (i & i) return ret ans = 0 pres = [0] * (len(nums) + 1) b = [0] * (len(nums) + 2) for i in range(0, len(nums)): pres[i + 1] = pres[i] + nums[i] sortedPres = sorted(pres) for end in pres: count = sumRange(b, bisect.bisect_right(sortedPres, end  lower))  sumRange(b, bisect.bisect_left(sortedPres, end  upper)) ans += count update(b, bisect.bisect_left(sortedPres, end) + 1, 1) return ans