天天看點

【LeetCode-二分查找】尋找兩個正序數組的中位數

題目描述

給定兩個大小為 m 和 n 的正序(從小到大)數組 nums1 和 nums2。

請你找出這兩個正序數組的中位數,并且要求算法的時間複雜度為 O(log(m + n))。

你可以假設 nums1 和 nums2 不會同時為空。

示例:

nums1 = [1, 3]
nums2 = [2]
則中位數是 2.0

nums1 = [1, 2]
nums2 = [3, 4]
則中位數是 (2 + 3)/2 = 2.5
           

題目連結: https://leetcode-cn.com/problems/median-of-two-sorted-arrays/

思路

因為題目要求算法的時間複雜度為 O(log(m + n)),log 的時間複雜度一般是使用二分查找來做。

因為兩個數組是有序的,是以我們可以一半一半地對數組進行尋找。假設兩個數組如下

【LeetCode-二分查找】尋找兩個正序數組的中位數

圖來自這篇題解。

假設我們要找第 k = 7 小的數字,則我們可以比較 nums1[k/2] 和 nums2[k/2] 之間的大小,如果 nums2[k/2]<nums1[k/2],因為數組是升序的,則 nums2[k/2] 之前的數字都不會是第 k 小的數字,可以直接删除

【LeetCode-二分查找】尋找兩個正序數組的中位數

黃色部分表示已删除。此時我們已經删除了 7/2 = 3 個數字,是以接下來我們要找第 k = 7 - 3 = 4 小的數字。同樣地,我們繼續比較 nums1[k/2] 和 nums2[k/2] 之間的大小,根據大小關系判斷如何删除元素。

這個過程會有一個問題,因為我們每次都是比較兩個數組的第 k/2 個元素,但可能會出現數組長度小于 k/2 的情況,如下

【LeetCode-二分查找】尋找兩個正序數組的中位數

解決辦法就是我們取 min(k/2, len) - 1 對應的元素進行比較,len 為數組的剩餘長度。

代碼如下:

class Solution {
public:
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        int m = nums1.size();
        int n = nums2.size();
        int left = (m+n+1)/2;
        int right = (m+n+2)/2;
        return (find(nums1, 0, m-1, nums2, 0, n-1, left) + find(nums1, 0, m-1, nums2, 0, n-1, right)) / 2;
    }

    double find(vector<int>& nums1, int left1, int right1, vector<int>& nums2, int left2, int right2, int k){
        int len1 = right1-left1+1;
        int len2 = right2-left2+1;
        if(len1>len2) return find(nums2, left2, right2, nums1, left1, right1, k); // 保證 len1 不大于 len2
        if(len1==0) return nums2[left2+k-1];  // 如果 len1==0,直接傳回目前 nums2 的第 k 個元素
        if(k==1) return min(nums1[left1], nums2[left2]);  // 如果 k==1,則傳回兩個數組頭更小的那個元素

        int i = left1 + min(k/2, len1)-1;  // 第 k/2 個元素,為了防止數組長度小于 k/2,用了 min 函數 
        int j = left2 + min(k/2, len2)-1;  // 同上
        if(nums1[i]<nums2[j])
            return find(nums1, i+1, right1, nums2, left2, right2, k-(i-left1+1)); // 别忘了最後一個參數
        else 
            return find(nums1, left1, right1, nums2, j+1, right2, k-(j-left2+1));
    }
};
           

代碼中還有一點,就是我們同時計算了第 left 小的元素和第 right 小的元素,然後求兩者的平均值作為答案。這樣做的原因是,假設兩個數組長度分别為 len1 和 len2,如果 len1+len2 為偶數,則我們需要求中間兩個數的平均值,中間兩個數就是第 left 小的元素和第 right 小的元素;如果 len1+len2 為奇數,則中間的那個數就是中位數,這時第 left 小的元素和第 right 小的元素是一樣的,會發生重複計算。

參考

https://leetcode-cn.com/problems/median-of-two-sorted-arrays/solution/xiang-xi-tong-su-de-si-lu-fen-xi-duo-jie-fa-by-w-2/

繼續閱讀