又是一道有意思的题目, Count of Range Sum 。(PS:leetcode 我已经做了 190 道,欢迎围观全部题解 https://github.com/hanzichi/leetcode )
题意非常简单,给一个数组,如果该数组的一个子数组,元素之和大于等于给定的一个参数值(lower),小于等于一个给定的参数值(upper),那么这为一组解,求总共有几组解。
一个非常容易想到的解法是两层 for 循环遍历子数组首尾,加起来判断,时间复杂度 O(n^2)。
/** * @param {number[]} nums * @param {number} lower * @param {number} upper * @return {number} */ var countRangeSum = function(nums, lower, upper) { var len = nums.length; var ans = 0; for (var i = 0; i < len; i++) { var sum = 0; for (var j = i; j < len; j++) { sum += nums[j]; if (sum >= lower && sum <= upper) ans++; } } return ans; };
交了下 TLE 了,看了下测试数据,数组长度为 9000,复杂度达到了 8100w,还是蛮大的。其实题目中也说了: A naive algorithm of O(n2) is trivial. You MUST do better than that.
如何将复杂度降到 log 级别?想到了二分的方法。可以将子数组和转换成两个前缀数组和的差,定义数组 sum, sum[i] 表示数组前 i 个元素的和,特殊的, sum[0]=0,那么元素 i 到元素 j 的和可以表示为 sum[j]-sum[i-1]。我们枚举 0 到 nums.length,比如枚举到了 sum[j],我们需要求满足条件的 i(i<j),sum[j]-sum[i] 的值满足大于等于 lower,小于等于 upper。我们需要枚举 sum[0] 到 sum[i],复杂度还是 O(n^2),如果 sum[0] 到 sum[i] 有序呢?
解法似乎呼之而出,用二分维护有序数组(用 splice 插入),同时用二分找到临界的数据,一次迭代需要多次二分。二分查找相关可以看我以前的文章 二分查找大集合(妈妈再也不用担心我的二分查找了) 。
注意下二分的边界,代码很容易写出来。
function binarySearch1(a, target) { target += 1; var start = 0 , end = a.length - 1; while(start <= end) { var mid = ~~((start + end) >> 1); if (a[mid] >= target) end = mid - 1; else start = mid + 1; } return start; } function binarySearch2(a, target) { var start = 0 , end = a.length - 1; while(start <= end) { var mid = ~~((start + end) >> 1); if (a[mid] >= target) end = mid - 1; else start = mid + 1; } return end; } var countRangeSum = function(nums, lower, upper) { var len = nums.length; var sum = []; var ans = 0; var num = 0; sum.push(0); for (var i = 0; i < len; i++) { ans += nums[i]; var a = ans - upper; var b = ans - lower; var pos1 = binarySearch2(sum, a) + 1; var pos2 = binarySearch1(sum, b) - 1; num += pos2 - pos1 + 1; var pos3 = binarySearch1(sum, ans); sum.splice(pos3, 0, ans); } return num; };
很不幸,还是 TLE 了,究其原因,我觉得应该是调用了 n 次 splice 方法。 感觉维护一棵二叉搜索树应该是可行的,无奈不会手写二叉搜索树 = =
那么可行的解法是什么呢?答案是归并排序的 "另类使用"。这里不讲归并排序,关于归并排序,可见我以前的文章 http://www.cnblogs.com/zichi/p/4796727.html 。
言归正传,首先预处理数组的前缀和,保存到数组 sum 中。然后用归并排序对数组 sum 进行排序,归并排序中有一步调用 merge 函数,将有序的左数组和右数组进行合并, 而这时的右数组中的任一元素在 sum 数组中的位置正是在左数组任一元素之后 !利用这,我们可以在 merge 前,对 left 数组和 right 数组满足条件的元素进行求解。
这个函数我定义为 getAns:
// 返回 b[j] - a[i] 值在 [wlower, wupper] 范围内组数 function getAns(a, b) { var sum = 0; var lena = a.length; var lenb = b.length; var start = 0; var end = 0; for (var i = 0; i < lenb; i++) { // to get start while (b[i] - a[start] >= wlower) { start++; } // to get end while (b[i] - a[end] > wupper) { end++; } sum += start - end; } return sum; }
做完一次归并排序,每次 left 和 right 数组合并前进行判断,就将所有 sum[j]-sum[i](j>i) 的情况进行了判断,简直神奇!
完整代码参考我的 Github https://github.com/hanzichi/leetcode/blob/master/Algorithms/Count%20of%20Range%20Sum/count-of-range-sum.js
224ms!Your runtime beats 100.00% of javascriptsubmissions 还是有点小激动