[LintCode/LeetCode] Median of two Sorted Arrays

359 查看

Problem

There are two sorted arrays A and B of size m and n respectively. Find the median of the two sorted arrays.

Example

Given A=[1,2,3,4,5,6] and B=[2,3,4,5], the median is 3.5.

Given A=[1,2,3] and B=[4,5], the median is 3.

Challenge

The overall run time complexity should be O(log (m+n)).

Note

由于要求O(log(m+n))的时间,所以选择二分法。
思路是找到两个数组合并起来的第k个元素。这样,只需计算两个数组的中位数是第几个元素,代入功能函数即可。当然,这里要先判断总长度的奇偶性,若为奇数,则取中间的元素;若为偶数,则计算中间两个数的平均值。

有两点需要注意:
在功能函数findKth()中,虽然是寻找第k个元素,但是对应数组的顺序是k-1,
由于二分搜索需要不断递归,所以极限情况是k1的时候,此时要返回AB中首元素较小的那个。

二分搜索第k个元素的原理:
首先,找到k个元素的中点mid,令mid = k/2-1
假设数组A和B的长度都大于首元素加上mid个元素的长度,那么,对于我们的数组A的前mid个元素和数组B的前mid个元素,一定不包含要找的第k个元素。
据此,根据二分法的性质,我们在递归时可以将前mid+1(即k/2)个元素排除。
那么,排除A还是B的前mid+1个元素呢?
比较一下A和B的第mid+1个元素,哪个小就排除该数组的前mid+1个元素。
然后,继续递归查找第k-k/2个元素。

Solution

    class Solution {
        public double findMedianSortedArrays(int[] A, int[] B) {
            int len = A.length + B.length;
            if (len % 2 == 0) return (findKth(A, 0, B, 0, len/2)+findKth(A, 0, B, 0, len/2+1))/2.0;
            else return findKth(A, 0, B, 0, len/2+1);
        }
        public double findKth(int[] A, int Astart, int[] B, int Bstart, int k) {
            if (Astart == A.length) return B[Bstart+k-1];
            else if (Bstart == B.length) return A[Astart+k-1];
            if (k == 1) return Math.min(A[Astart], B[Bstart]);
            int mid = k/2-1, kNew = k-k/2, kA = Integer.MAX_VALUE, kB = kA;
            if (Astart + mid < A.length) kA = A[Astart+mid];
            if (Bstart + mid < B.length) kB = B[Bstart+mid];
            if (kA < kB) return findKth(A, Astart+k/2, B, Bstart, kNew);
            else return findKth(A, Astart, B, Bstart+k/2, kNew);
        }
    }

Naive method: O(m+n) time

public class Solution {
    public double findMedianSortedArrays(int[] nums1, int[] nums2) {
        int n1 = nums1.length, n2 = nums2.length;
        int len = n1+n2;
        if (len%2 == 0) return (find(nums1, nums2, len/2)+find(nums1, nums2, len/2+1))/2.0;
        else return (double)find(nums1, nums2, len/2+1);
    }
    public int find(int[] A, int[] B, int k) {
        int res = 0;
        int i = 0, j = 0;
        while (k != 0 && i < A.length && j < B.length) {
            if (A[i] < B[j]) res = A[i++];
            else res = B[j++];
            k--;
        }
        if (k != 0) {
            if (i < A.length) res = A[i+k-1];
            else res = B[j+k-1];
        }
        return res;
    }
}