327 Count of Range Sum

1. Question

Given an integer arraynums, return the number of range sums that lie in[lower, upper]inclusive. Range sumS(i, j)is defined as the sum of the elements innumsbetween indicesiandj(ij), inclusive.
Note: A naive algorithm of O(n^2) is trivial. You MUST do better than that.
Example: Givennums=[-2, 5, -1],lower=-2,upper=2, Return3. The three ranges are :[0, 0],[2, 2],[0, 2]and their respective sums are:-2, -1, 2.

2. Implementation

(1) TreeMap
我们定义sum[i]为nums[0...i]之间的subarray sum, 对于区间[j,i]的subarray sum,其值等于sum[i] - sum[j], 如果 lower <= sum[i] - sum[j] <= upper, 通过推导可以得出 sum[j] - upper <= sum[j] <= sum[i] - lower. 我们遍历数组的时候,利用sum变量加上当前的数,这样我们很容易得到sum[i], 要得到sum[j], 我们利用TreeMap的区间查询的特性,查询treemap,key在 sum[i] - upper[i] 和 sum[i] - lower之间的值的个数,最后得到结果
1
public class Solution {
2
public int countRangeSum(int[] nums, int lower, int upper) {
3
TreeMap<Long, Long> map = new TreeMap();
4
map.put((long)0, (long)1);
5
long sum = 0;
6
int count = 0;
7
for (int i = 0; i < nums.length; i++) {
8
sum += nums[i];
9
long fromKey = sum - upper;
10
long toKey = sum - lower;
11
Map<Long, Long> subMap = map.subMap(fromKey, true, toKey, true);
12
13
for (long value : subMap.values()) {
14
count += value;
15
}
16
17
if (map.containsKey(sum)) {
18
map.put(sum, map.get(sum) + 1);
19
}
20
else {
21
map.put(sum, (long)1);
22
}
23
}
24
return count;
25
}
26
}
Copied!
(2) Merge Sort
思路: 这里的思路和Leetcode 315很类似,我们利用merge sort,在merge的过程中count 符合条件的区间和个数. 具体在我们的merge()里,我们知道sum[start, mid] 和 sum[mid + 1, end]这两个区间是已经排好序的,对于每个在左区间的数sum[left], 我们分别要找出一个index high满足sum[high] - sum[left] > upper, 以及第一个index low满足sum[low] - sum[left] >= lower, 这样符合条件的区间个数就是 high - low. 在计算符合条件的区间个数同时,我们也顺便完成merge的过程
1
class Solution {
2
public int countRangeSum(int[] nums, int lower, int upper) {
3
if (nums == null || nums.length == 0) {
4
return 0;
5
}
6
7
int n = nums.length;
8
long[] sum = new long[n + 1];
9
10
for (int i = 1; i <= n; i++) {
11
sum[i] = sum[i - 1] + nums[i - 1];
12
}
13
14
int[] count = new int[1];
15
mergeSort(sum, 0, sum.length - 1, lower, upper, count);
16
return count[0];
17
}
18
19
public void mergeSort(long[] sum, int start, int end, int lower, int upper, int[] count) {
20
if (start >= end) {
21
return;
22
}
23
24
int mid = start + (end - start) / 2;
25
mergeSort(sum, start, mid, lower, upper, count);
26
mergeSort(sum, mid + 1, end, lower, upper, count);
27
merge(sum, start, mid, end, lower, upper, count);
28
}
29
30
public void merge(long[] sum, int start, int mid, int end, int lower, int upper, int[] count) {
31
long[] temp = new long[end - start + 1];
32
33
int low = mid + 1;
34
int high = mid + 1;
35
int right = mid + 1;
36
int tIndex = 0;
37
38
for (int left = start; left <= mid; left++) {
39
while (low <= end && sum[low] - sum[left] < lower) ++low;
40
while (high <= end && sum[high] - sum[left] <= upper) ++high;
41
42
count[0] += high - low;
43
44
while (right <= end && sum[right] <= sum[left]) temp[tIndex++] = sum[right++];
45
temp[tIndex++] = sum[left];
46
}
47
48
while (right <= end) {
49
temp[tIndex++] = sum[right++];
50
}
51
52
for (int k = start; k <= end; k++) {
53
sum[k] = temp[k - start];
54
}
55
}
56
}
Copied!
(3) Segment Tree
1
Copied!

3. Time & Space Complexity

TreeMap: 时间复杂度O(nlogn), 空间复杂度O(n)
Merge Sort: 时间复杂度O(nlogn), 空间复杂度O(n)