4. Median of Two Sorted Arrays

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 1:

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

The median is 2.0

Example 2:

nums1 = [1, 2]
nums2 = [3, 4]

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

思路

定义函数getKthNumber(int A[], int l1, int r1, int B[], int l2, int r2, int k),表示在第一个数组区间l1, r1和第二个数组区间l2, r2内寻找第k小的数，我们找到当前A,B数组中间的数A[m1], B[m2],并统计出[l1, m1], [l2, m2] 区间的数的个数cnt， 如果A[m1] < B[m2]且cnt>k, 那么答案不可能在[m2, r2]这个区间，如果A[m1]<B[m2]且cnt<=k,那么我们丢弃[l1, m1]这个区间。A[m1] > B[m2]同理。 A[m1] == B[m2]可随意规划为上面两个状态。

class Solution {
public:
    int getKthNumber(vector<int>& A, int l1, int r1, vector<int>& B, int l2, int r2, int k) {
       if(l1 > r1)
          return B[l2+k-1];
       if(l2 > r2)
           return A[l1+k-1];

       int m1 = (l1 + r1) / 2;
       int m2 = (l2 + r2) / 2;
       int cnt = m1 - l1 + m2 - l2 + 2;

       if(A[m1] < B[m2]) {
           if(cnt > k)
                return getKthNumber(A, l1, r1, B, l2, m2-1, k);
           else
               return getKthNumber(A, m1+1, r1, B, l2, r2, k-(m1-l1+1));
        }else{
           if(cnt > k)
               return getKthNumber(A, l1, m1-1, B, l2, r2, k);
           else
               return getKthNumber(A, l1, r1, B, m2+1, r2, k-(m2-l2+1));
       }
    }
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        if((nums1.size()+nums2.size())%2 == 1) {
            return getKthNumber(nums1, 0, nums1.size()-1, nums2, 0, nums2.size()-1, (nums1.size()+nums2.size())/2+1);
        }else{
            double a1 = getKthNumber(nums1, 0, nums1.size()-1, nums2, 0, nums2.size()-1, (nums1.size()+nums2.size())/2);
            double a2 = getKthNumber(nums1, 0, nums1.size()-1, nums2, 0, nums2.size()-1, (nums1.size()+nums2.size())/2+1);
            return (a1 + a2) / 2.0;
        }
    }
};