算法之美：你可能想不到的归并排序的神奇应用

580 查看

又是一道有意思的题目，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

解法似乎呼之而出，用二分维护有序数组（用 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 方法。感觉维护一棵二叉搜索树应该是可行的，无奈不会手写二叉搜索树 = =

那么可行的解法是什么呢？答案是归并排序的 “另类使用”。这里不讲归并排序，关于归并排序，可见我以前的文章。

言归正传，首先预处理数组的前缀和，保存到数组 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 。

224ms！Your runtime beats 100.00% of javascript submissions

还是有点小激动