327. Count of Range Sum (Hard)
https://leetcode.com/problems/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(n2) 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
.
Solutions
public class Solution {
public int countRangeSum(int[] nums, int lower, int upper) {
if (nums == null || nums.length == 0) {
return 0;
}
int ans = 0;
// sums of all preceding elements, use long in case overflow.
long[] sums = new long[nums.length + 1];
Arrays.fill(sums, 0);
// keep the count of same sum, elements in TreeMap is sorted by key increasingly.
TreeMap<Long, Integer> sumMap = new TreeMap<>();
sumMap.put(0L, 1);
for (int i = 0; i < nums.length; i++) {
sums[i + 1] = sums[i] + nums[i];
if (!sumMap.containsKey(sums[i + 1])) {
sumMap.put(sums[i + 1], 0);
}
sumMap.put(sums[i + 1], sumMap.get(sums[i + 1]) + 1);
}
for (int i = 0; i < sums.length; i++) {
// Find out all the ranges that start from i (inclusive) that are subject to lower<= rangeSum <=upper.
// Given that, sums before i should be removed
sumMap.put(sums[i], sumMap.get(sums[i]) - 1);
if (sumMap.get(sums[i]) == 0) {
sumMap.remove(sums[i]);
}
// range sum starting from i to x is sum[x] - sum[i], subject to lower < sum[x] - sum[i] < upper
Map<Long, Integer> subMap = sumMap.subMap(sums[i] + lower, sums[i] + upper + 1);
for (Long key: subMap.keySet()) {
ans += subMap.get(key);
}
}
return ans;
}
}