# 4 - Median of Two Sorted Arrays

最直觀的解法就是把兩個 array 花 O(m+n) 的時間 merge 起來,再用 O(1) 找到中位數就好,不過 we can do better!

我們可以先把這個問題想成找第 k 大的 element,然後找中位數就只是其中的一種 k 而已。接下來的重點是,我們要一次捨棄 k/2 個 element,而做法就是

  1. 比較兩個 array 裡面第 k/2 大的 element

  2. 如果 A array 第 k/2 大的 element 比 B array 第 k/2 大的 element 小,那這 A 的前 k/2 個 element 都可以丟掉,因為如果把兩個 array merge 起來,A 的前 k/2 個 element 一定不可能是第 k 大的。

  3. 重新搜尋第 k-k/2 大的 element。直到 k 是 1 或某一個 array 的 element 已經被用完。

class Solution {
public:
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        int length = nums1.size() + nums2.size();
        if(length%2 == 0)
            return (findKth(nums1, 0, nums2, 0, length/2)+findKth(nums1, 0, nums2, 0, length/2+1))/2.0;
        else
            return  findKth(nums1, 0, nums2, 0, length/2+1);
    }

private: 
    int findKth(vector<int>& A, int A_start, vector<int>& B, int B_start, int k) {
        if(A_start >= A.size()) 
            return B[B_start+k-1];
        if(B_start >= B.size())
            return A[A_start+k-1];

        if(k == 1)
            return min(A[A_start], B[B_start]);

        int A_val = A_start+k/2-1 < A.size() ? A[A_start+k/2-1] : INT_MAX;
        int B_val = B_start+k/2-1 < B.size() ? B[B_start+k/2-1] : INT_MAX;

        if(A_val < B_val)
            return findKth(A, A_start+k/2, B, B_start, k-k/2);
        else
            return findKth(A, A_start, B, B_start+k/2, k-k/2);
    }
};

Last updated