LeetCode 327 - Count of Range Sum

作者 QIFAN 日期 2017-04-26
LeetCode 327 - Count of Range Sum

原题链接:https://leetcode.com/problems/count-of-range-sum/#/description


题干

给定一个数组,返回数字之和落在 [lower, upper] 内的区间的个数。

举例:
输入: nums = [-2, 5, -1], lower = -2, upper = 2
输出: 3
解释: 这三个区间范围分别是 [0,0],对应到数组是 [-2] ,和为 -2 ,满足;[2,2] ,对应到数组是 [-1] ,和为 -1 ,满足;[0,2] ,对应到数组是 [-2, 5, -1] 和为 2 ,满足。

思路

识别子问题。对于每一个区间 [lo, hi],已知其和为 sum,下一步的区间情况:

  • hi++ ,lo 不变 ,sum += nums[hi]
  • hi++ ,lo++ , sum += nums[hi] - nums[lo - 1]
  • hi++ ,lo– , sum += nums[hi] + nums[lo]

然后再用一个 memo 矩阵存已经算过的范围的和,大致思路就出来了。这样子时间复杂度是 $O(N^2)$ ,只是省去了重复计算和。

记得处理 integer overflow 问题。

代码:

int count = 0;
boolean[][] memo;
int N;
public int countRangeSum(int[] nums, int lower, int upper) {
if (nums == null || nums.length == 0) {
return 0;
}
N = nums.length;
memo = new boolean[N][N];
calculateRangeSum(0, 0, nums[0], nums, lower, upper);
return count;
}
private void calculateRangeSum(int lo, int hi, long sum, int[] nums, int lower, int upper) {
if (lo > hi || memo[lo][hi]) {
return;
}
if (sum >= lower && sum <= upper) {
count++;
}
if (hi < N - 1 && !memo[lo][hi + 1]) {
if (lo < N - 1 && !memo[lo + 1][hi + 1]) {
calculateRangeSum(lo + 1, hi + 1, sum + nums[hi + 1] - nums[lo], nums, lower, upper);
}
if (lo > 0 && !memo[lo - 1][hi + 1]) {
calculateRangeSum(lo - 1, hi + 1, sum + nums[hi + 1] + nums[lo - 1], nums, lower, upper);
}
calculateRangeSum(lo, hi + 1, sum + nums[hi + 1], nums, lower, upper);
}
memo[lo][hi] = true;
}

这个方法还是超时了。于是看了 discuss 的 binary search 做法,发现是另一种分解子问题的思路。

https://discuss.leetcode.com/topic/34108/summary-of-the-divide-and-conquer-based-and-binary-indexed-tree-based-solutions/4

对于 (0, n-1) index 的数组范围,设满足的范围数为 T(0, n)
T(0, n) = T(0, m) + T(m + 1, n - 1) + C
其中 m = (n - 1) / 2 ,问题就转为了 (0, m) 范围,与 (m+1, n-1) 范围之和,与部分在前半部分,部分在后半部分的满足范围数,把这个新问题设为 C 。

乍一看这个新问题 C 与原问题没有区别,粗暴解法还是 $O(N^2)$ 。有一个 trick 的地方就在于将范围分成了前后两部分,所以只要算出了前面的部分的和都装进一个数组,后面部分的和也装进数组,并对后面部分进行排序,那样就可以用二叉搜索 nlogn 的复杂度找到满足的范围。所以总的算法复杂度是 $O(N(logN)^2)$ 。

原链接中还提了好些优化的方法,包括 Binary Index Tree ,可以点进去看。另外上一篇介绍 BIT 的题:LeetCode 308 - Range Sum Query 2D - Mutable
代码如下:

public int countRangeSum(int[] nums, int lower, int upper) {
if (nums == null || nums.length == 0 || lower > upper) return 0;
return countRangeSumSub(nums, 0, nums.length - 1, lower, upper);
}
private int countRangeSumSub(int[] nums, int l, int r, int lower, int upper) {
if (l == r) return nums[l] >= lower && nums[r] <= upper ? 1 : 0; // base case
int m = l + (r - l) / 2;
long[] arr = new long[r - m]; // prefix array for the second subarray
long sum = 0;
int count = 0;
for (int i = m + 1; i <= r; i++) {
sum += nums[i];
arr[i - (m + 1)] = sum; // compute the prefix array
}
Arrays.sort(arr); // sort the prefix array
// Here we can compute the suffix array element by element.
// For each element in the suffix array, we compute the corresponding
// "insertion" indices of the modified bounds in the sorted prefix array
// then the number of valid ranges sums will be given by the indices difference.
// I modified the bounds to be "double" to avoid duplicate elements.
sum = 0;
for (int i = m; i >= l; i--) {
sum += nums[i];
count += findIndex(arr, upper - sum + 0.5) - findIndex(arr, lower - sum - 0.5);
}
return countRangeSumSub(nums, l, m, lower, upper) + countRangeSumSub(nums, m + 1, r, lower, upper) + count;
}
// binary search function
private int findIndex(long[] arr, double val) {
int l = 0, r = arr.length - 1, m = 0;
while (l <= r) {
m = l + (r - l) / 2;
if (arr[m] <= val) {
l = m + 1;
} else {
r = m - 1;
}
}
return l;
}