【算法之美】你可能想不到的歸並排序的神奇應用 — leetcode 327. Count of Range Sum

又是一道有意思的題目，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 of Range Sum/count-of-range-sum.js

224ms！Your runtime beats 100.00% of javascriptsubmissions 還是有點小激動