題目描述
給定兩個大小為 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 的時間復雜度一般是使用二分查找來做。
因為兩個數組是有序的,所以我們可以一半一半地對數組進行尋找。假設兩個數組如下
圖來自這篇題解。
假設我們要找第 k = 7 小的數字,則我們可以比較 nums1[k/2] 和 nums2[k/2] 之間的大小,如果 nums2[k/2]<nums1[k/2],因為數組是升序的,則 nums2[k/2] 之前的數字都不會是第 k 小的數字,可以直接刪除
黃色部分表示已刪除。此時我們已經刪除了 7/2 = 3 個數字,所以接下來我們要找第 k = 7 - 3 = 4 小的數字。同樣地,我們繼續比較 nums1[k/2] 和 nums2[k/2] 之間的大小,根據大小關系判斷如何刪除元素。
這個過程會有一個問題,因為我們每次都是比較兩個數組的第 k/2 個元素,但可能會出現數組長度小於 k/2 的情況,如下
解決辦法就是我們取 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 小的元素是一樣的,會發生重復計算。