又是一道有意思的題目,Count of Range Sum。(PS:leetcode 我已經做了 190 道,歡迎圍觀全部題解 https://github.com/hanzichi/leetcode)
題意非常簡單,給一個陣列,如果該陣列的一個子陣列,元素之和大於等於給定的一個引數值(lower),小於等於一個給定的引數值(upper),那麼這為一組解,求總共有幾組解。
一個非常容易想到的解法是兩層 for 迴圈遍歷子陣列首尾,加起來判斷,時間複雜度 O(n^2)。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
/** * @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 插入),同時用二分找到臨界的資料,一次迭代需要多次二分。二分查詢相關可以看我以前的文章 二分查詢大集合(媽媽再也不用擔心我的二分查詢了)。
注意下二分的邊界,程式碼很容易寫出來。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
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:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 |
/ 返回 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
還是有點小激動
打賞支援我寫出更多好文章,謝謝!
打賞作者
打賞支援我寫出更多好文章,謝謝!