Wednesday, January 7, 2015

LeetCode 4: Median of Two Sorted Arrays

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)).
public class Solution {
    public double findMedianSortedArrays(int A[], int B[]) {
        // This is a special binary search
        // We should respectively consider when (m+n) is odd or even number.
        int m = A.length;
        int n = B.length;
        int len = m+n;
        
        // If len is odd number
        if (len%2 == 1)
            return findKth(A, 0, m, B, 0, n, len/2);
        // If len is even number
        else
            return (findKth(A, 0, m, B, 0, n, len/2-1)+findKth(A, 0, m, B, 0, n, len/2))/2.0;
    }
    
    // Find the kth number in A+B
    // Note: k is 0-based.
    private double findKth(int A[], int loA, int hiA, int B[], int loB, int hiB, int k)
    {
        int lenA = hiA-loA;
        int lenB = hiB-loB;
        
        // If A is empty, return the kth number of B.
        if (lenA == 0)
            return B[loB+k];
        
        // If B is empty, return the kth number of A.    
        if (lenB == 0)
            return A[loA+k];
        
        // If k == 0, the kth number is the first number of A+B.    
        if (k == 0)
            return Math.min(A[loA], B[loB]);
            
        int midA = loA + (hiA-loA)/2;
        int midB = loB + (hiB-loB)/2;
        
        // A is divided into A1 and A2 by A[midA];
        // B is divided into B1 and B2 by B[midB].
        if (A[midA] <= B[midB])
        {
            // If true, abandon B2.
            if (k <= lenA/2+lenB/2)
                return findKth(A, loA, hiA, B, loB, midB, k);
            // Abandon A1, update k.
            else
                return findKth(A, midA+1, hiA, B, loB, hiB, k-lenA/2-1);
     }
     else
     {
         // If true, abandon A2.
            if (k <= lenA/2+lenB/2)
                return findKth(A, loA, midA, B, loB, hiB, k);
            // Abandon B1, update k.
            else
                return findKth(A, loA, hiA, B, midB+1, hiB, k-lenB/2-1);
     }
    }
}

No comments:

Post a Comment