Ecommet

Median of Two Sorted Arrays

Hard Problem 47 of 151 Java and Python

Problem

Given two sorted arrays nums1 and nums2 of sizes m and n, return the median of the two arrays combined.

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

Example 1

Input
nums1 = [1, 3], nums2 = [2]
Output
2.0
Why
Merged: [1, 2, 3].

Example 2

Input
nums1 = [1, 2], nums2 = [3, 4]
Output
2.5
Why
Merged: [1, 2, 3, 4], median (2 + 3) / 2.

Constraints

  • 0 ≤ m, n ≤ 1000, 1 ≤ m + n ≤ 2000
  • -10⁶ ≤ nums1[i], nums2[i] ≤ 10⁶

Explanation

The median splits the merged array into a left half and a right half of equal size. Take i elements from nums1 and half - i from nums2 for the left half. The split is right when every left element is ≤ every right element: nums1[i - 1] ≤ nums2[j] and nums2[j - 1] ≤ nums1[i].

If nums1[i - 1] is too big, take fewer from nums1; if nums2[j - 1] is too big, take more. That is a binary search on i, done over the smaller array.

Approach 1 · Naive

Merge, then pick the middle

  • Time O(m + n)
  • Space O(m + n)

Merge the two sorted arrays like in merge sort and read the middle element(s).

  1. Merge into one sorted array with two pointers.
  2. Return the middle element, or the average of the two middle elements.
Java
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 (j == nums2.length || (i < nums1.length && nums1[i] <= nums2[j])) {
                merged[k++] = nums1[i++];
            } else {
                merged[k++] = nums2[j++];
            }
        }
        int n = merged.length;
        return n % 2 == 1 ? merged[n / 2] : (merged[n / 2 - 1] + merged[n / 2]) / 2.0;
    }
}
Python
class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        merged = []
        i = j = 0
        while i < len(nums1) or j < len(nums2):
            if j == len(nums2) or (i < len(nums1) and nums1[i] <= nums2[j]):
                merged.append(nums1[i])
                i += 1
            else:
                merged.append(nums2[j])
                j += 1
        n = len(merged)
        return merged[n // 2] if n % 2 else (merged[n // 2 - 1] + merged[n // 2]) / 2

Approach 2 · Optimal

Binary search the partition

  • Time O(log(min(m, n)))
  • Space O(1)

Binary-search how many elements of the smaller array go into the left half, and fix the rest from the other array.

  1. Make nums1 the shorter array. Let half = (m + n + 1) / 2.
  2. Search i in [0, m]; j = half - i. Read L1, R1 around i in nums1 and L2, R2 around j in nums2, using ±∞ past the ends.
  3. If L1 <= R2 and L2 <= R1, the split is right: median = max(L1, L2) if the total is odd, else (max(L1, L2) + min(R1, R2)) / 2.
  4. If L1 > R2, move i left; otherwise move it right.
Java
class Solution {
    public double findMedianSortedArrays(int[] nums1, int[] nums2) {
        if (nums1.length > nums2.length) {
            return findMedianSortedArrays(nums2, nums1);
        }
        int m = nums1.length;
        int n = nums2.length;
        int half = (m + n + 1) / 2;
        int lo = 0;
        int hi = m;
        while (lo <= hi) {
            int i = (lo + hi) / 2;
            int j = half - i;
            int left1 = i == 0 ? Integer.MIN_VALUE : nums1[i - 1];
            int right1 = i == m ? Integer.MAX_VALUE : nums1[i];
            int left2 = j == 0 ? Integer.MIN_VALUE : nums2[j - 1];
            int right2 = j == n ? Integer.MAX_VALUE : nums2[j];
            if (left1 <= right2 && left2 <= right1) {
                if ((m + n) % 2 == 1) {
                    return Math.max(left1, left2);
                }
                return (Math.max(left1, left2) + (double) Math.min(right1, right2)) / 2.0;
            } else if (left1 > right2) {
                hi = i - 1;
            } else {
                lo = i + 1;
            }
        }
        return 0;
    }
}
Python
class Solution:
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
        if len(nums1) > len(nums2):
            nums1, nums2 = nums2, nums1
        m, n = len(nums1), len(nums2)
        half = (m + n + 1) // 2
        lo, hi = 0, m
        while lo <= hi:
            i = (lo + hi) // 2
            j = half - i
            left1 = nums1[i - 1] if i > 0 else float("-inf")
            right1 = nums1[i] if i < m else float("inf")
            left2 = nums2[j - 1] if j > 0 else float("-inf")
            right2 = nums2[j] if j < n else float("inf")
            if left1 <= right2 and left2 <= right1:
                if (m + n) % 2:
                    return max(left1, left2)
                return (max(left1, left2) + min(right1, right2)) / 2
            if left1 > right2:
                hi = i - 1
            else:
                lo = i + 1
        return 0.0

Complexity compared

ApproachTimeSpace
Naive Merge, then pick the middleO(m + n)O(m + n)
Optimal Binary search the partitionO(log(min(m, n)))O(1)

Key takeaways

  • Search over the shorter array so j never goes out of range.
  • Infinity sentinels at the edges remove a pile of special cases.

Your notes

Saved in this browser as you type.

Solved it on your own?

Mark it done to track your progress, or star it to come back before an interview.

Next problem