4. Median of Two Sorted Arrays

NeetCode: Median of Two Sorted Arrays · LeetCode: 4. Median of Two Sorted Arrays · Difficulty: Hard

Given two sorted arrays nums1 and nums2, return the median of all their elements combined. The required time is O(log(m + n)). The result is a double.

Example

Input:  nums1 = [1, 2], nums2 = [3]       Output: 2.0
Input:  nums1 = [1, 3], nums2 = [2, 4]    Output: 2.5

Brute force: merge and take the middle

Merge the arrays like the merge step of merge sort, then read the median from the merged array: the middle element for odd length, the average of the two middle elements for even length. It is correct but runs in O(m + n), which does not meet the required logarithmic bound.

The sum of two middle int values can exceed Integer.MAX_VALUE, so the code adds them as long before dividing by 2.0.

class Solution {
    public double findMedianSortedArrays(int[] nums1, int[] nums2) {
        int[] merged = new int[nums1.length + nums2.length];
        int i = 0;
        int j = 0;
        int k = 0;
        while (i < nums1.length && j < nums2.length) {
            if (nums1[i] <= nums2[j]) {
                merged[k++] = nums1[i++];
            } else {
                merged[k++] = nums2[j++];
            }
        }
        while (i < nums1.length) {
            merged[k++] = nums1[i++];
        }
        while (j < nums2.length) {
            merged[k++] = nums2[j++];
        }

        int n = merged.length;
        if (n % 2 == 1) {
            return merged[n / 2];
        }
        return ((long) merged[n / 2 - 1] + merged[n / 2]) / 2.0;
    }
}

public class Main {
    public static void main(String[] args) {
        Solution s = new Solution();
        System.out.println(s.findMedianSortedArrays(new int[] {1, 2}, new int[] {3}));
        System.out.println(s.findMedianSortedArrays(new int[] {1, 3}, new int[] {2, 4}));
    }
}

2.0
2.5

Time: O(m + n). Space: O(m + n).

Optimal: binary search the partition

Idea. The median depends only on where the "left half" ends. Cut a after i elements and b after j elements, choosing i + j = half (half of all elements, rounded up). Name the four values touching the cuts:

a:  ... left1 | right1 ...
b:  ... left2 | right2 ...

Invariant. The cut is correct exactly when left1 <= right2 and left2 <= right1 (each array is already sorted, so those are the only two checks). Then everything on the left is <= everything on the right, and the median comes from max(left1, left2) (and min(right1, right2) when the total is even).

Search. Binary search i on the shorter array (so j is never negative); j = half - i follows. If left1 > right2, too many elements were taken from a, so move i left. If left2 > right1, too few were taken, so move i right.

Sentinels. If a cut is at an array's edge there is no element there: use Integer.MIN_VALUE for a missing left value and Integer.MAX_VALUE for a missing right value, so the comparisons still hold.

class Solution {
    public double findMedianSortedArrays(int[] nums1, int[] nums2) {
        // Always binary search the shorter array.
        int[] a = nums1.length <= nums2.length ? nums1 : nums2;
        int[] b = nums1.length <= nums2.length ? nums2 : nums1;
        int m = a.length;
        int n = b.length;
        int half = (m + n + 1) / 2;

        int lo = 0;
        int hi = m;
        while (lo <= hi) {
            int i = lo + (hi - lo) / 2; // elements of a on the left
            int j = half - i;           // elements of b on the left

            int left1 = i == 0 ? Integer.MIN_VALUE : a[i - 1];
            int right1 = i == m ? Integer.MAX_VALUE : a[i];
            int left2 = j == 0 ? Integer.MIN_VALUE : b[j - 1];
            int right2 = j == n ? Integer.MAX_VALUE : b[j];

            if (left1 > right2) {
                hi = i - 1; // took too many from a
            } else if (left2 > right1) {
                lo = i + 1; // took too few from a
            } else {
                int leftMax = Math.max(left1, left2);
                if ((m + n) % 2 == 1) {
                    return leftMax;
                }
                int rightMin = Math.min(right1, right2);
                return ((long) leftMax + rightMin) / 2.0;
            }
        }
        throw new IllegalArgumentException("inputs must be sorted");
    }
}

public class Main {
    public static void main(String[] args) {
        Solution s = new Solution();
        System.out.println(s.findMedianSortedArrays(new int[] {1, 2}, new int[] {3}));
        System.out.println(s.findMedianSortedArrays(new int[] {1, 3}, new int[] {2, 4}));
        System.out.println(s.findMedianSortedArrays(new int[] {}, new int[] {5}));
        System.out.println(s.findMedianSortedArrays(new int[] {Integer.MAX_VALUE}, new int[] {Integer.MAX_VALUE}));
    }
}

2.0
2.5
5.0
2.147483647E9

Walk-through for a = [1, 3], b = [2, 4]: half = 2. Try i = 1, j = 1: left1 = 1, right1 = 3, left2 = 2, right2 = 4. Both checks pass, total is even, so the answer is (max(1, 2) + min(3, 4)) / 2 = 2.5.

Time: O(log(min(m, n))). Space: O(1).

hi = i - 1 cannot go below -1, and the loop then ends: left1 > right2 is impossible when i == 0 because left1 is then Integer.MIN_VALUE. The last call in main has two values of Integer.MAX_VALUE. The sum is computed as long, so the average is 2147483647 (Java prints it as 2.147483647E9) and does not wrap.

How to explain it in an interview

• Baseline: merge the two arrays and read the middle, O(m + n) time and space.
• Key idea: the merged array is not needed, only a cut that splits all elements into a left half and a right half where every left value is <= every right value. That cut can be binary searched.
• Steps: (1) pick a as the shorter array; (2) binary search i, set j = half - i; (3) read left1/right1/left2/right2 with MIN_VALUE/MAX_VALUE sentinels at the edges; (4) if left1 > right2 go left, if left2 > right1 go right; (5) otherwise compute the median from max(left1, left2) (and min(right1, right2) if even).
• Complexity: O(log(min(m, n))) time, O(1) space.
• Common questions: searching the shorter array keeps j = half - i in range and gives the smaller logarithm. Empty arrays are handled by the sentinels. The average is computed in long and divided by 2.0 so that the sum cannot overflow int.