4. Median of Two Sorted Arrays

There are two sorted arrays nums1 and nums2 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)).

Example 1:

nums1 = [1, 3]
nums2 = [2]

The median is 2.0

Example 2:

nums1 = [1, 2]
nums2 = [3, 4]

The median is (2 + 3)/2 = 2.5

Solution:

If the total length of two sorted arrays is odd. the median is one number from the array. Otherwise, the median is the average of two numbers from the arrays.

We are actually trying to find the kth smallest number in two sorted arrays.

To find it, we check the (k / 2)-th smallest number (let's say valueA )in array A and (k / 2)-th smallest number (let's say valueB) in array B.

If valueA < valueB, we know we can safely remove all the numbers in array A which are smaller than valueA (including valueA itself). And reduce the problem to find (k - k / 2)-th smallest number in two sorted arrays.

If valueB > valueA, we know we can safely remove all the numbers in array B which are smaller than valueB (including valueB itself).

If we find either the remaining array A or B's length is smaller than k / 2, we know we cannot remove element in it. Thus we remove k / 2 elements from the other array.

The time complexity is O(log(m + n)).



Code:

public class Solution {
    public double findMedianSortedArrays(int[] nums1, int[] nums2) {
        int len = nums1.length + nums2.length;
        if (len % 2 == 0) {
            return (findKthElement(nums1, 0, nums2, 0, len / 2) + findKthElement(nums1, 0, nums2, 0, len / 2 + 1)) / 2.0;
        }
        return findKthElement(nums1, 0, nums2, 0, len / 2 + 1);
    }
    
    public int findKthElement(int[] A, int startA, int[] B, int startB, int k) {
        if (startA >= A.length) {
            return B[startB + k - 1];
        }
        if (startB >= B.length) {
            return A[startA + k - 1];
        }
        if (k == 1) {
            return Math.min(A[startA], B[startB]);
        }
        
        int valueInA = startA + k / 2 - 1 < A.length ? A[startA + k / 2 - 1] : Integer.MAX_VALUE;
        int valueInB = startB + k / 2 - 1 < B.length ? B[startB + k / 2 - 1] : Integer.MAX_VALUE;
        
        if (valueInA < valueInB) {
            return findKthElement(A, startA + k / 2, B, startB, k - k / 2);
        }
        return findKthElement(A, startA, B, startB + k / 2, k - k / 2);
    }
}