Wednesday, April 18, 2012

Find K-th element in two sorted arrays

[Question]:
Find the k-th element in two sorted arrays.

[Analysis]:
The k-th element exists in A[0, ... , k-1] and B[0, ..., k-1]. Suppose m of the k elements are in A,  then A[m]> B[k-m-1], B[k-m]>A[m-1]. The k-th element is MAX (A[m-1], B[k-m-1]). Use binary search to find such m. Time complexity is O(log k).

[Alternative Question]:
Find the median of two sorted arrays. In this case, k = (M+N)/2, M, N are lengths of A and B.

[Solution]:
//
//Find the k-th element.
//
    public static int findKthElementOfArrays(int[] a, int[] b, int k) {
        assert k>0;
        assert a.length + b.length >= k;
     
        int low = Math.max(0, k-b.length);
        int high= Math.min(k, a.length);
        int m = 0;
     
        while (low < high) {
            m = low + (high-low)/2;
         
            if (k-m-1>=0 && a[m]<b[k-m-1]) {
                low = m+1;
            }
            else if (m>0 && k-m <b.length && b[k-m]<a[m-1]) {
                high = m;
            }
            else {
                low = m;
                break;
            }
        }
        if (low == 0) return b[k-1];
        if (low == k) return a[k-1];
        return Math.max(a[low-1], b[k-low-1]);
    }
//
// Find the median.
//
class Solution {
public:
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        vector<int> &a=nums1.size()<nums2.size()?nums1:nums2,
                &b = a==nums1?nums2:nums1;
       
        auto median=[](vector<int>& x) {
            int i = x.size()/2;
            return x.size()%2?(double)x[i]: (x[i-1]+x[i])/2.0;
        };
        if (a.empty()) return median(b);
        if (b.empty()) return median(a);
       
        int m=a.size(), n=b.size(), k=(m+n)>>1;
        int lo=max(0,k-n), hi=min(m, k);
        while (lo<hi) {
            int mid = (lo+hi)>>1;
            if (k-mid-1>=0 && b[k-mid-1]>a[mid])
                lo = mid+1;
            else if (mid-1>=0 && b[k-mid]<a[mid-1])
                hi = mid;
            else {
                lo = mid;
                break;
            }
        }
     
        int upper = (lo==m)?b[k-lo]:(k-lo)==n?a[lo]:min(a[lo],b[k-lo]);
        int lower = (lo==0)?b[k-1]:(lo==k)?a[k-1]:max(a[lo-1],b[k-lo-1]);
       
        return (m+n)&0x01? upper: (upper+lower)/2.0;
    }
};

No comments:

Post a Comment