LeetCode 4. Median of Two Sorted Arrays

Question

Given two sorted arrays nums1 and nums2 of size m and n respectively, return the median of the two sorted arrays.

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

Example 1:

Input: nums1 = [1,3], nums2 = [2]
Output: 2.00000
Explanation: merged array = [1,2,3] and median is 2.

Example 2:

Input: nums1 = [1,2], nums2 = [3,4]
Output: 2.50000
Explanation: merged array = [1,2,3,4] and median is (2 + 3) / 2 = 2.5.

Constraints:

• nums1.length == m
• nums2.length == n
• 0 <= m <= 1000
• 0 <= n <= 1000
• 1 <= m + n <= 2000
• -106 <= nums1[i], nums2[i] <= 106

Source: https://leetcode.com/problems/median-of-two-sorted-arrays/

Solution

A simple solution is to maintain two pointers for two sorted arrays, and use linear search to find the median. Obviously, it is simple but slow.

Because both arrays are sorted, we can use binary search to improve performance. Remember that the essence of binary search is to exclude a certain percentage of candiates from the search range on each round.

// Time complexity: O(log(m+n)), because k = (m+n)/2
public double findMedianSortedArrays(int[] nums1, int[] nums2) {
    int m = nums1.length;
    int n = nums2.length;
    int lmedian = (m + n + 1) / 2;
    int rmedian = (m + n + 2) / 2;
    if (lmedian == rmedian) {
        return getKth(nums1, nums2, 0, 0, lmedian);
    } else {
        return (getKth(nums1, nums2, 0, 0, lmedian) +
                getKth(nums1, nums2, 0, 0, rmedian)) / 2.0;
    }
}

// get the k-th smallest (starts from 1) element in two sorted arrays
// start is the start index of search window (inclusive)
// @pre: total number of elements in these two arrays is no less than k
// Time complexity: O(log(k))
private static int getKth(int[] nums1, int[] nums2, int start1, int start2, int k) {
    int m = nums1.length;
    int n = nums2.length;
    if (start1 > m - 1) { // used up nums1
        return nums2[start2 + k - 1];
    }
    if (start2 > n - 1) { // used up nums2
        return nums1[start1 + k - 1];
    }
    // to make sure k/2 > 0
    if (k == 1) {
        return Math.min(nums1[start1], nums2[start2]);
    }

    // the k/2 th element from start
    // Finding the kth element is a recursive process
    // So using rounded down k/2 is safe
    int kmid1 = Integer.MAX_VALUE;
    int kmid2 = Integer.MAX_VALUE;
    if (start1 + k / 2 - 1 < m) { // nums1 has k/2 th element from start
        kmid1 = nums1[start1 + k / 2 - 1];
    }
    if (start2 + k / 2 - 1 < n) { // nums2 has k/2 th element from start
        kmid2 = nums2[start2 + k / 2 - 1];
    }

    // k/2 elements on the smaller side must be included in the smallest k elements
    // remove them from further searches
    if (kmid1 < kmid2) {
        return getKth(nums1, nums2, start1 + k / 2, start2, k - k / 2);
    } else {
        // if two mids are equal, we can remove either k/2 elements
        return getKth(nums1, nums2, start1, start2 + k / 2, k - k / 2);
    }
}