[LeetCode] Median of Two Sorted Arrays

题目

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)).

思路

题目中已经给出了时间复杂度的要求，所以不能采用暴力求解的方法，我第一个思路是想到了一种O(m+n)的解法，代码如下：

public class Solution {
    public double findMedianSortedArrays(int[] nums1, int[] nums2) {
        int i = 0,j = 0, k = 0;
        int m = nums1.length, n = nums2.length;
        int[] mergeNums = new int[m + n];
        
        while(i < m && j < n){
            if(nums1[i] <= nums2[j])
                mergeNums[k++] = nums1[i++];
            else 
                mergeNums[k++] = nums2[j++];
        }
        
        while(k < m + n){
            if(i < m)
                mergeNums[k++] = nums1[i++];
            else
                mergeNums[k++] = nums2[j++];
        }
        
        if((m + n) % 2 == 0)
            return ((double)mergeNums[(m + n) / 2] + (double)mergeNums[(m + n - 2) / 2]) / 2;
        else
            return mergeNums[(m + n - 1) / 2]; 
    }
}

尽管这个答案在LeetCode OJ上就能AC了，但是我还是想寻求更优的解法。如果合并两个数组后排序，显然即使采用快速排序，时间复杂度也还有(m+n)log(m+n)。因为题目中给出的时间复杂度是O(log(m+n))，印象中只有二分查找法的时间复杂度是这样，所以有可能需要将二分查找的思想应用到这道题里面，即每次可以去除当前数组范围内一半的元素。

具体思路是，首先将问题一般化为在数组A和B中找到第k小的数，做这样的假设，数组A和数组B的元素个数都大于k/2，同时不考虑奇偶性的问题，这两点在具体的代码实现的时候会考虑。题目中已经给出数组是排好序的，所以A[k-1/2]和B[k-1/2]分别表示A和B的下标为k-1/2的元素，即有k-1/2个数小于这个元素。比较这两个数，一共有三种情况：

1.当A[k-1/2]=B[k-1/2]时，这个相等的元素就是要找的数，原因是在A和B中分别有k-1/2个元素小于m，所以m是第k小的数，当k=(m+n)/2时即为中位数。

2.当A[k-1/2]<B[k-1/2]时，这表明A[k-1/2]不可能大于两个数组合并之后的第k小值，所以我们可以将A数组中该下标前的元素抛弃。

3.当A[k-1/2]>B[k-1/2]时，与2类似。

最终代码

public class Solution {
    public double findMedianSortedArrays(int[] nums1, int[] nums2) {
        int m = nums1.length, n = nums2.length;
        
        //if m + n is odd
        if ((m + n) % 2 == 1)
            return findKth(nums1, 0, m, nums2, 0, n, (m + n) / 2 + 1);  
        else
            return  (findKth(nums1, 0, m, nums2, 0, n, (m + n) / 2) + findKth(nums1, 0, m, nums2, 0, n, (m + n) / 2 + 1)) / 2;
    }
    
    public double findKth(int A[], int l1, int r1, int B[], int l2, int r2, int k) {  
        if (r1 > r2) 
            return findKth(B, l2, r2, A, l1, r1, k);  
        if (r1 == 0)   
            return B[l2 + k - 1];  
        if (k == 1)   
            return Math.min(A[l1], B[l2]);  
  
        int p1 = Math.min(k / 2, r1), p2 = k - p1;  
      
        if (A[l1 + p1 - 1] < B[l2 + p2 - 1])   
            return findKth(A, l1 + p1, r1 - p1, B, l2, r2, k - p1);  
        else   
            return findKth(A, l1, r1, B, l2 + p2, r2 - p2, k - p2);  
    }  
}