Leetcode: 4. Median of Two Sorted Arrays

Description

There are two sorted arrays nums1 and nums2 of size m and n respectively.

Find the median of the two sorted arrays. The overall run time complexity should be O(log (m+n)).

Example

Example1:
    nums1 = [1, 3]
    nums2 = [2]

    The median is 2.0
    
Example2:
    nums1 = [1, 2]
    nums2 = [3, 4]

    The median is (2 + 3)/2 = 2.5

思路

• 每次做这个题都要纠结半天。好好理理吧。
• 此题的另一种通用形式为：给定两个已经排好序的数组，找到两者所有元素中第K大的元素
  • 分析1：O(m+n), 直接merge, 然后找出第k大的数。 可以部分优化，即i = 0, pa = 0, pb = 0, 然后比较，若 A[pa] < B[pb],则i++, pa++。同理。。然后直接到第K大元素即可。当K接近(m+n)时，优化不大
  • 分析2：O(log(m + n))， 二分
    • 假设A和B中的元素个数都大于k/2，我们将A的第k/2个元素(即A[k/2 - 1])和B的第k/2个元素相比较，可以得到：
      • A[k/2 - 1] == B[k/2 - 1]
      • A[k/2 - 1] < B[k/2 - 1]
      • A[k/2 - 1] > B[k/2 - 1]
    • 如果A[k/2 - 1] 小于B[k/2 - 1],这说明A[0]到A[k/2 - 1]这部分元素肯定在top k范围内，所以可以删除A的这k/2个元素。同理A[k/2 - 1] > B[k/2 - 1]时，可以删除B的这k/2个元素
    • 当A[k/2 - 1] == B[k/2 - 1]时，说明找到第K大的元素，直接返回A[k/2-1]或B[k/2-1]
    • 因此，我们可以写一个递归函数。那么函数什么时候应该终止呢？
      • 当 A 或 B 是空时，直接返回 B[k-1] 或 A[k-1]；
      • 当 k=1 是，返回 min(A[0], B[0])；
      • 当 A[k/2-1] == B[k/2-1] 时，返回 A[k/2-1] 或 B[k/2-1]

代码

• 时间复杂度 O(log(m+n))

class Solution {
public:
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        int len1 = nums1.size();
        int len2 = nums2.size();
        
        if(len1 == 0 && len2 == 0)
            return 0;
        else if(len1 == 0){
            return len2 % 2 == 0 ? ((nums2[len2/2 - 1] + nums2[len2/2])/2.0 ): nums2[len2 / 2]; 
        }else if(len2 == 0){
             return len1 % 2 == 0  ? ((nums1[len1/2 - 1] + nums1[len1/2]) / 2.0) : nums1[len1 / 2]; 
        } 
        
        int total = len1 + len2;
        
        if(total & 0x1)
            return find_kth(&(*nums1.begin()), len1, &(*nums2.begin()), len2, total / 2 + 1);
        else{
            return  (find_kth(&(*nums1.begin()), len1, &(*nums2.begin()), len2, total / 2)
               + find_kth(&(*nums1.begin()), len1, &(*nums2.begin()), len2, total / 2 + 1)) / 2.0;
        }
    }
    
private:
    int find_kth(int *A, int m, int *B, int n, int k){
        if(m > n) return find_kth(B, n, A, m, k);
        if(m == 0) return B[k - 1];
        if(k == 1) return min(A[0], B[0]);
        
        int ia = min(k / 2, m), ib = k - ia;
        if(A[ia - 1] < B[ib - 1])
            return find_kth(A + ia, m - ia, B, n, k - ia);
        else if(A[ia - 1] > B[ib - 1])
            return find_kth(A, m, B + ib, n - ib, k - ib);
        return A[ia - 1];
    }
};