Thursday, December 11, 2014

Median of Two Sorted Arrays

Q:
There are two sorted arrays A and B 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)).

T:
For 0<=i<=m and 0<=j<=n, we split A and B into two parts. The lower part is A[0:i-1] and B[0:j-1], the upper part is A[i:m-1] and B[j:n-1]. If i and j split the two part equally, i.e. i+j = m-i + n-j, and all the elements in lower part is smaller than the upper part, then the median can be calculated from A[i-1], A[i], B[j-1], B[j]. Note in the code below j is calculated as (m+n+1)/2 - i, so that when m+n is odd, i+j == (m+n)/2 +1.

A:

class Solution {
public:
    double findMedianSortedArrays(int A[], int m, int B[], int n) {
        if(m < n) {
            return findMedianSortedArrays(B, n, A, m);
        }
        if(m==0) {
            return 0;
        }
        if(n==0) {
            if(m%2==1) {
                return A[(m-1)/2];
            } else {
                return (A[(m-1)/2] + A[(m-1)/2+1])/2.0;
            }
        }
        int l=0;
        int r=m;
        while(l<=r) {
            int i=(l+r)/2;
            int j=(m+n+1)/2-i;
            if(i < m && j>0 && A[i] < B[j-1]) {
                l = i+1;
            } else if(i>0 && j < n && A[i-1] > B[j]) {
                r = i-1;
            } else {
                int num1 = 0;
                if(i==0 && j>0) {
                    num1 = B[j-1];
                } else if(j==0 && i>0) {
                    num1 = A[i-1];
                } else if(i>0 && j>0) {
                    num1 = max(A[i-1], B[j-1]);
                }
                if ((m+n)%2 == 1) {
                    return num1;
                }
                int num2 = 0;
                if (i==m && j < n) {
                    num2 = B[j];
                } else if(j==n && i < m) {
                    num2 = A[i];
                } else if(i < m && j < n) {
                    num2 = min(A[i], B[j]);
                }
                return ((double)num1 + (double)num2) / 2.0;
            }
        }
        return 0;
    }
};

No comments:

Post a Comment