This page explains Java solution to problem Count of Range Sum
using TreeMap
data structure.
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 1: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.
package com.vc.hard;
import java.util.*;
class CountOfRangeSum {
public int countRangeSum(int[] nums, int lower, int upper) {
TreeMap<Long, Integer> prefixSum = new TreeMap<>();
long sum = 0;
int count = 0;
prefixSum.put(0l, 1);
for(int i = 0; i < nums.length; i++) {
sum += nums[i];
long from = sum - upper;
long to = sum - lower;
Map<Long, Integer> subMap = prefixSum.subMap(from, true, to, true);
for(int value: subMap.values()) count += value;
prefixSum.put(sum, prefixSum.getOrDefault(sum, 0) + 1);
}
return count;
}
}
O(N * log(N)) Where
N is total number of elements in an input array
O(N) Where
N is total number of elements in an input array