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.

Challenge:

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

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.

Solution:

In order to solve this problem, we should find the k-th smallest number from both A and B. Then we could find the middle number for odd length and middle two numbers for even length to calculate the median.

A: 1, 2, 3, 4, 5, 6, inf, inf

B: 2, 3, 4, 5, inf, inf

We are going to find the k=6 now:

1. divide k by 2 is 3, we would like to drop 3 numbers from A or B

2. compare the third number in A and B, A[2] = 3 and B[2] = 4. We are going to drop first three number from A to have

   A: 4, 5, 6, inf, inf

   B: 2, 3, 4, 5, inf, inf

3. since we dropped three numbers from A, our k is 3 now (6 - 3).

4. then we repeat step 1 to divide k by 2 again to have 1. Now, we will drop 1 number from A or B.

5. compare the first number in A and B, A[0] = 4 and B[0] = 2, we will drop 2 in B now.

   A: 4, 5, 6, inf, inf

   B: 3, 4, 5, inf, inf

6. our new k is 2 (3 - 1). After divide by 2, we would like to drop 1 number again by same procedure. Then, we have

   A: 4, 5, 6, inf, inf

   B: 4, 5, inf, inf

7. Finally k is 1, we choose the smallest number from first element of A or B as the k-th smallest number. In fact, this number is the median for A and B.

Code:

class Solution {
    /**
     * @param A: An integer array.
     * @param B: An integer array.
     * @return: a double whose format is *.5 or *.0
     */
    public double findMedianSortedArrays(int[] A, int[] B) {
        if (A == null || B == null) {
            return 0;
        }
        int n = A.length + B.length;
        int k = n / 2;
        double median = findKthNumber(A, B, 0, 0, k + 1);
        if ((n & 1) == 0) {
            median += findKthNumber(A, B, 0, 0, k);
            median = median / 2;
        }
        return median;
    }

    private double findKthNumber(int[] A, int[] B, int aidx, int bidx, int k) {
        if (aidx >= A.length) {
            return B[k - 1 + bidx];
        }
        if (bidx >= B.length) {
            return A[k - 1 + aidx];
        }
        if (k == 1) {
            if (A[aidx] < B[bidx]) {
                return A[aidx];
            } else {
                return B[bidx];
            }
        }

int keyA = Integer.MAX_VALUE;
        if (aidx + k / 2 - 1 < A.length) {
            keyA = A[aidx + k / 2 - 1];
        }
        int keyB = Integer.MAX_VALUE;
        if (bidx + k / 2 - 1 < B.length) {
            keyB = B[bidx + k / 2 - 1];
        }
        if (keyA <= keyB) {
            return findKthNumber(A, B, aidx + k / 2, bidx, k - k /2);
        } else {
            return findKthNumber(A, B, aidx, bidx + k / 2, k - k / 2);
        }
    }
}