Median of Array

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

Find the median of the two sorted arrays ( The median of the array formed by merging both the arrays ).

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

Sample Input

A : [1 4 5]
B : [2 3]

Sample Output

NOTE: IF the number of elements in the merged array is even, then the median is the average of n / 2 th and n/2 + 1th element.
For example, if the array is [1 2 3 4], the median is (2 + 3) / 2.0 = 2.5

思路：

We make the following observation:
For array A and B:
A : [1 4 5]
i
B : [2 3 6]
j-1 j
if Ai > Bj-1 && Ai < Bj, then Ai is the i + j + 1 th element
similiary if Bj > Ai - 1 && Bj < Ai, then Bj is the i + j + 1 th element
In the example, 4 < 3 and 4 < 6, hence 4 is the 1 + 2 + 1 = 4th element.
This is because we can know that 4 is greater than 1 element in its own array, and 2 elements in the other array.

Based on this observation, we can do a binary search to find the kth largest element in the two sorted array.
We initialize i and j based on the value of k. and check whether Ai or Bj meets the requirement, if so then it is the kth element.

k = 3

A : [3 4 5]
i
B : [1 2 6]
j

Else if Ai > Bj, then we know Ai is too large to be the kth element, it is at least the k + 1 th element, so we can ignore all elements after it including itself.
Also, Because Ai > Bj, we know Bj <= Ai - 1, otherwise Bj would have passed the test. Because Bj <= Ai - 1, it is too small to be the kth element. (it could be in the case Bj == Ai -1, but we can still ignore Bj since Ai - 1 will still be included in the next search), so we can ignore Bj and elements before it.
Hence, we can search for (k - j - 1) th element, in A[0, i) and B[j + 1, n).

We can do similar analysis on Bj > Ai. In the case when Bj == Ai, both are the kth element, just return either of them. (since Ai/Bj >= j + i elements)

e.g
k = 3

A : [3 4 5]
i
B : [1 4 6]
j

Solution:

Time: O(log m + n)
Space: O(1)

public class Solution {
// DO NOT MODIFY BOTH THE LISTS
public double findMedianSortedArrays(final List<Integer> a, final List<Integer> b) {
int m = a.size();
int n = b.size();
int mid = (m + n) / 2;
if ((m + n) % 2 == 0) {
// 1 2
// 3 4
return (kth(a, b, mid) + kth(a, b, mid + 1)) / 2.0;
} else {
// 1 2
// 3
return kth(a, b, mid + 1);
}
}

private int kth(List<Integer> a, List<Integer> b, int k) {
/*
A : [1 4 5]
i
B : [2 3 6]
j-1 j
obersavation:
if Ai > Bj-1 && Ai < Bj, then Ai is the i + j + 1 th element
similiary if Bj < Ai - 1 && Bj < Ai, then Bj is the i + j + 1 th element
*/
int m = a.size();
int n = b.size();
if (k == 1) {
if (a.size() == 0) return b.get(0);
else if (b.size() == 0) return a.get(0);
else {
return a.get(0) <= b.get(0) ? a.get(0) : b.get(0);
}
}
int i = (int) ((m + 0.0) / (m + n) * (k - 1)); // make sure i < m
int j = k - 1 - i;
int Ai = i == m ? Integer.MAX_VALUE : a.get(i);
int Bj = j == n ? Integer.MAX_VALUE : b.get(j);
int Ai_1 = i == 0 ? Integer.MIN_VALUE : a.get(i - 1);
int Bj_1 = j == 0 ? Integer.MIN_VALUE : b.get(j - 1);
if (Ai > Bj_1 && Ai < Bj) {
return Ai;
}
if (Bj > Ai_1 && Bj < Ai) {
return Bj;
}
if (Ai > Bj) {
// can ignore [i + 1, m) and [0, j]
// we can reduce k by j + 1, because [0, j] are the smaller numbers getting eliminated
return kth(a.subList(0, i), b.subList(j + 1, n), k - j - 1);
} else if (Bj > Ai){
return kth(a.subList(i + 1, m), b.subList(0, j), k - i - 1);
} else {
return Ai;
}
}
}

class Solution {
// A: [1 4 5]
// i
// B: [2 5 6]
// j
public double findMedianSortedArrays(int[] nums1, int[] nums2) {
int m = nums1.length;
int n = nums2.length;
int len = m + n;
if (len % 2 == 0) {
return (kth(nums1, nums2, len / 2) + kth(nums1, nums2, len / 2 + 1)) / 2.0;
} else {
return kth(nums1, nums2, len / 2 + 1);
}
}

private int kth(int[] nums1, int[] nums2, int k) {
int m = nums1.length;
int n = nums2.length;
int i = (int) ((m + 0.0) / (m + n) * (k - 1));
int j = k - 1 - i;
int Ai = i == m ? Integer.MAX_VALUE : nums1[i];
int Ai_1 = i == 0 ? Integer.MIN_VALUE : nums1[i - 1];
int Bj = j == n ? Integer.MAX_VALUE : nums2[j];
int Bj_1 = j == 0 ? Integer.MIN_VALUE : nums2[j - 1];
if (Bj_1 < Ai && Ai < Bj) {
return Ai;
}
if (Ai_1 < Bj && Bj < Ai) {
return Bj;
}
if (Ai < Bj) {
return kth(Arrays.copyOfRange(nums1, i + 1, m), Arrays.copyOfRange(nums2, 0, j), k - i - 1);
} else if (Ai > Bj) {
return kth(Arrays.copyOfRange(nums1, 0, i), Arrays.copyOfRange(nums2, j + 1, n), k - j - 1);
} else {
return Ai;
}
}
}

class Solution {
    // A: [1 4 5]
    //       i
    // B: [2 5 6]
    //       j
    public double findMedianSortedArrays(int[] nums1, int[] nums2) {
        int m = nums1.length;
        int n = nums2.length;
        int len = m + n;
        if (len % 2 == 0) {
            return (kth(nums1, nums2, len / 2) + kth(nums1, nums2, len / 2 + 1)) / 2.0;
        } else {
            return kth(nums1, nums2, len / 2 + 1);
        }
    }
    
    private int kth(int[] nums1, int[] nums2, int k) {
        int m = nums1.length;
        int n = nums2.length;
        int i = (int) ((m + 0.0) / (m + n) * (k - 1));
        int j = k - 1 - i;
        int Ai = i == m ? Integer.MAX_VALUE : nums1[i];
        int Ai_1 = i == 0 ? Integer.MIN_VALUE : nums1[i - 1];
        int Bj = j == n ? Integer.MAX_VALUE : nums2[j];
        int Bj_1 = j == 0 ? Integer.MIN_VALUE : nums2[j - 1];
        if (Bj_1 <= Ai && Ai <= Bj) {
            return Ai;
        }
        if (Ai_1 <= Bj && Bj <= Ai) {
            return Bj;
        }
        if (Ai < Bj) {
            return kth(Arrays.copyOfRange(nums1, i + 1, m), Arrays.copyOfRange(nums2, 0, j), k - i - 1);
        } else {
            return kth(Arrays.copyOfRange(nums1, 0, i), Arrays.copyOfRange(nums2, j + 1, n), k - j - 1);
        }
    }
}