問題描述
給定一個整數(shù)數(shù)組念秧，返回range sum 落在給定區(qū)間[lower, upper] (包含lower和upper)的個數(shù)缚态。range sum S(i, j) 表示數(shù)組中第i 個元素到j 個元素之和夏跷。
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.
分析
這個題目比較難，樓主第一次面對這種題型，直接繳械投降。參考了各位大神的解題思路匣缘，總結了兩種解法。一種是TreeMap思路鲜棠，另外一種是使用segment tree (or binary index tree)肌厨。題目尋找需要range sum 在[lower, upper] 之間的個數(shù)，滿足條件的case用數(shù)學公式表達為：
lower <= sum[i] - sum[j] <= upper, i > j, sum[i] is prefix sum of nums at index of i.
也就是
sum[i] - high <= sum[j] <= sum[i] - lower, i > j, sum[i] is prefix sum of nums at index of i.
(or
lower + sum[j] <= sum[i] <= sum[j] + higher, i > j, sum[i] is prefix sum of nums at index of i.)
那么我們的問題可以轉化為求落在[sum[i] - high,sum[i] - lower] 區(qū)間sum[j]的個數(shù), i = 0....n, j < i岔留。
無論是TreeMap還是Segment Tree夏哭，總體的時間復雜度都為nlogn。
實現(xiàn)
TreeMap
TreeMap 的key 是prefixsum, value 是相對應的個數(shù)献联。主要使用TreeMap的subMap的方法竖配，求得落在區(qū)間內[sum[i] - high, sum[i] - lower]的sum[j]的個數(shù)。
public int countRangeSum(int[] nums, int lower, int upper) {
if(nums == null || nums.length == 0){
return 0;
}
//key is the sum[i], value is the corresponding count
// sum[i] - sum[j] in [lower, upper], transform to find how many sum[j] 在區(qū)間[sum[i] - high, sum[i] - lower]里逆。
TreeMap<Long, Integer> map = new TreeMap();
long sum = 0;
int cnt = 0;
for(int i = 0; i < nums.length; i++){
sum += nums[i];
//sum[0, i]滿足case
if(sum >= lower && sum <= upper){
cnt++;
}
//find sum[j] 的個數(shù)that lies in [sum[i] - high, sum[i] - lower]之間
cnt += map.subMap(sum - upper, true, sum - lower, true).values().stream().mapToInt(Integer::valueOf).sum();
map.put(sum, map.getOrDefault(sum, 0) + 1);
}
return cnt;
}
Segment Tree
Segment Tree每個節(jié)點保存區(qū)間段的范圍和落在這個區(qū)間內prefix sum的個數(shù)进胯。
class Node {
Node left;
Node right;
//落在區(qū)間內的個數(shù)
int count;
long min;
long max;
public Node(long min, long max) {
this.min = min;
this.max = max;
}
}
//構建segement tree
private Node buildTree(Long[] valArr, int low, int high) {
if(low > high) return null;
Node root = new Node(valArr[low], valArr[high]);
if(low == high) return root;
int mid = low + (high - low)/2;
root.left = buildTree(valArr, low, mid);
root.right = buildTree(valArr, mid+1, high);
return root;
}
private void update(Node root, Long val) {
if(root == null) return;
if(val >= root.min && val <= root.max) {
root.count++;
update(root.left, val);
update(root.right, val);
}
}
private int query(Node root, long min, long max) {
if(root == null) return 0;
if(min > root.max || max < root.min) return 0;
if(min <= root.min && max >= root.max) return root.count;
return query(root.left, min, max) + query(root.right, min, max);
}
public int countRangeSum(int[] nums, int lower, int upper) {
if(nums == null || nums.length == 0) return 0;
int ans = 0;
Set<Long> valSet = new HashSet<Long>();
long sum = 0;
for(int i = 0; i < nums.length; i++) {
sum += (long) nums[i];
valSet.add(sum);
}
Long[] valArr = valSet.toArray(new Long[0]);
Arrays.sort(valArr);
Node root = buildTree(valArr, 0, valArr.length-1);
sum = nums[0];
ans += (sum >= lower && sum <= upper) ? 1:0;
for(int i = 1; i < nums.length; i++) {
//sum[i]
update(root, sum);
//sum[j]
sum += (long) nums[i];
ans += (sum >= lower && sum <= upper) ? 1:0;
ans += query(root, (long)sum - upper, (long)sum - lower);
}
return ans;
}
