一刷 08/2022
Version #2 Merge Sort
看某一个prefix index后面有多少个prefix满足lower <= prefix[j]-prefix[i] <= upper
Time O(NlogN)
Space O(N)
class Solution {
public int countRangeSum(int[] nums, int lower, int upper) {
long[] prefixSum = new long[nums.length + 1];
for (int i = 0; i < nums.length; i++) {
prefixSum[i + 1] = prefixSum[i] + nums[i];
}
// find count of prefixSum on the right that lower <= prefixSum[j] - prefixSum[i] <= uppder
// [-2,5,-1,-4]
// [0, -2, 3, 2, -2]
// [-2, 0] [-2, 2, 3]
int[] count = new int[1];
mergeSort(prefixSum, new long[prefixSum.length], 0, prefixSum.length - 1, count, lower, upper);
return count[0];
}
private void mergeSort(long[] nums, long[] aux, int start, int end, int[] count, int lower, int upper) {
if (start >= end) {
return;
}
int mid = start + (end - start) / 2;
mergeSort(nums, aux, start, mid, count, lower, upper);
mergeSort(nums, aux, mid + 1, end, count, lower, upper);
// left is the first index that nums[left] - num[i] >= lower
// right is the first index that nums[right] - nums[i] > upper
int left = mid + 1, right = mid + 1;
// int i = mid + 1, j = mid + 1;
// for (int k = low; k <= mid; k++) {
// while (i <= high && pfxSum[i] - pfxSum[k] < lower) i++;
// while (j <= high && pfxSum[j] - pfxSum[k] <= upper) j++;
// count += j - i;
// }
// [0, -2, 3, 2]
// [-2, 0, 2, 3]
for (int i = start; i <= mid; i++) {
while (left <= end && nums[left] - nums[i] < lower) {
left++;
}
while (right <= end && nums[right] - nums[i] <= upper) {
right++;
}
count[0] += right - left;
}
for (int i = start; i <= end; i++) {
aux[i] = nums[i];
}
int pleft = start, pright = mid + 1;
int p = start;
while (pleft <= mid && pright <= end) {
if (aux[pleft] <= aux[pright]) {
nums[p++] = aux[pleft++];
} else {
nums[p++] = aux[pright++];
}
}
while (pleft <= mid) {
nums[p++] = aux[pleft++];
}
while (pright <= end) {
nums[p++] = aux[pright++];
}
}
}
Version #1 Segment Tree[TLE]
Time O(NlogN)
Space O(Range of all prefix sums)
class Solution {
class Node {
Node left, right;
int count;
int start, end;
public Node(int start, int end) {
this.start = start;
this.end = end;
}
}
public int countRangeSum(int[] nums, int lower, int upper) {
int[] prefixSum = new int[nums.length + 1];
int min = 0, max = 0;
for (int i = 0; i < nums.length; i++) {
prefixSum[i + 1] = prefixSum[i] + nums[i];
min = Math.min(min, prefixSum[i + 1]);
max = Math.max(max, prefixSum[i + 1]);
}
Node root = buildTree(min, max);
int count = 0;
for (int i = nums.length; i >= 0; i--) {
// prefixSum[i] - try to find the count of j that lower <= prefixSum[j] - prefixSum[i] <= upper
// lower + prefixSum[i] <= prefixSum[j] <= upper + prefixSum[j]
count += query(root, (long)lower + prefixSum[i], (long)upper + prefixSum[i]);
update(root, prefixSum[i]);
}
return count;
}
private void update(Node node, int num) {
if (node.start == num && node.end == num) {
node.count++;
return;
}
int mid = node.start + (node.end - node.start) / 2;
if (num <= mid) {
update(node.left, num);
} else {
update(node.right, num);
}
node.count = node.left.count + node.right.count;
}
private int query(Node node, long left, long right) {
// [start, end]
// [left, right]
if (left > node.end || right < node.start) {
return 0;
}
if (left <= node.start && right >= node.end) {
return node.count;
}
return query(node.left, left, right) + query(node.right, left, right);
}
private Node buildTree(int start, int end) {
Node n = new Node(start, end);
if (start == end) {
return n;
}
int mid = start + (end - start) / 2;
n.left = buildTree(start, mid);
n.right = buildTree(mid + 1, end);
return n;
}
}
