Median of Two Sorted Arrays

The problem

You are given two arrays of numbers, nums1 and nums2, each sorted from smallest to largest. Either one may be empty, but not both, and values may repeat. Return the median of all their values taken together: with the values in one sorted order, it is the middle one when their count is odd, and the average of the two middle ones when it is even. Aim for a method whose running time is O(log(m+n)), where m and n are the two lengths.

What to write

function solve(nums1: number[], nums2: number[]): number

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

Output
2

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

Output
2.5

Constraints

Aim for O(log(m+n)) time.

Hints

Each hint gives away a little more. Open one only when you need it.

1. Hint 1: Use the constraints

   The task asks for O(log(m+n)) time. Merging the two arrays, even only as far as the middle, handles about half of all the values, which is far more than that. So the median has to be found without building the merged order.

2. Hint 2: Ask a smaller question

   Take [1, 3, 8] and [2, 4, 9, 10]. Write the seven values in order. The median is the fourth, so mark the four smallest values. Where do those four sit in their own arrays?

3. Hint 3: Name the pattern

   Call the smaller half of all the values the left half. It is made of the first values of each array, so the only unknown is how many values nums1 gives it. If a count is too small, every smaller count is too. If it is too large, every larger count is too. That is binary search on the answer. Read the lesson on binary search on the answer.

4. Hint 4: State the invariant

   Cut each array where its part of the left half ends. The value before a cut is its left value, and the value after it is its right value, with minus or plus infinity where a cut is at an end. The cuts are correct exactly when each array's left value is at most the other array's right value.

5. Hint 5: Sketch the pseudocode

   Make nums1 the shorter array. Binary search how many values it gives, i, from 0 to its length; nums2 gives half the total, rounded up, minus i. Left value of nums1 above right value of nums2: lower i. The reverse: raise i. Otherwise return the largest left value, averaged with the smallest right value if the total is even.

Solution

From here on the page gives the method away. If you are still solving, go back to the hints.

Approach

Binary-search a partition of the shorter array so every value on the combined left is at most every value on the right.

The slow way

Merge the two arrays into one sorted list, then take its middle value, or the average of its two middle values when the count is even. It is correct, but merging handles every value, m + n of them, while the task asks for a number of steps that grows only like the logarithm of that.

The insight

Picture all the values in sorted order, split in the middle into a left half and a right half. The median sits at the split. Both arrays are sorted, so the left half is the first few values of nums1 plus the first few values of nums2. Its size is known, so choosing the count from nums1 fixes the count from nums2. Only that one count has to be found.

Step by step

1. Search the shorter array for how many values it gives the left half. The code first makes nums1 the shorter array. nums1 may then give the left half none of its values, all of them or anything between, and nums2 always has enough values for the rest.
2. Try the middle count; nums2 gives the rest of the left half. The left half holds half of all the values, rounded up, so it has the extra value when the total is odd. Taking the first values of each array cuts both in two. That pair of cuts is the partition.
3. Note the values next to each array's cut, or infinity past an end. The value before a cut is its left value, and the value after it is its right value. A cut at an end has no value on one side: minus or plus infinity stands in, so its comparison always passes.
4. The cuts are correct when no left value is above a right value. Each array is sorted, so its left value is never above its own right value. Only the two comparisons across the arrays can fail. When both hold, no left-half value is above any right-half value.
5. Read the median off the largest left value and the smallest right one. With an odd total the left half has one value more, so its largest value is the middle one. With an even total the two middle values are that one and the smallest right value, so average them.
6. Left value of nums1 above right value of nums2: take fewer from nums1. That left value belongs in the right half, so nums1 gave too many. Giving even more only moves its left value up and the right value of nums2 down, so every larger count fails as well.
7. Otherwise the left value of nums2 is too big: take more from nums1. Only one comparison is left to fail: the left value of nums2 is above the right value of nums1. So nums2 gave too many values and nums1 too few, and every smaller count fails the same way.

The same numbered steps are the comments in each solution.

function solve(nums1, nums2) {
  // 1. Search the shorter array for how many values it gives the left half
  if (nums1.length > nums2.length) return solve(nums2, nums1);
  const m = nums1.length,
    n = nums2.length;
  let l = 0,
    r = m;
  while (l <= r) {
    // 2. Try the middle count; nums2 gives the rest of the left half
    const i = (l + r) >>> 1,
      j = Math.floor((m + n + 1) / 2) - i;
    // 3. Note the values next to each array's cut, or infinity past an end
    const a = nums1[i - 1] ?? -Infinity,
      b = nums1[i] ?? Infinity,
      c = nums2[j - 1] ?? -Infinity,
      d = nums2[j] ?? Infinity;
    // 4. The cuts are correct when no left value is above a right value
    if (a <= d && c <= b) {
      // 5. Read the median off the largest left value and the smallest right one
      return (m + n) % 2
        ? Math.max(a, c)
        : (Math.max(a, c) + Math.min(b, d)) / 2;
    }
    // 6. Left value of nums1 above right value of nums2: take fewer from nums1
    if (a > d) {
      r = i - 1;
    } else {
      // 7. Otherwise the left value of nums2 is too big: take more from nums1
      l = i + 1;
    }
  }
  // Sorted arrays always have correct cuts, so this line is never reached.
  return 0;
}

function solve(nums1: number[], nums2: number[]): number {
  // 1. Search the shorter array for how many values it gives the left half
  if (nums1.length > nums2.length) return solve(nums2, nums1);
  const m = nums1.length,
    n = nums2.length;
  let l = 0,
    r = m;
  while (l <= r) {
    // 2. Try the middle count; nums2 gives the rest of the left half
    const i = (l + r) >>> 1,
      j = Math.floor((m + n + 1) / 2) - i;
    // 3. Note the values next to each array's cut, or infinity past an end
    const a = nums1[i - 1] ?? -Infinity,
      b = nums1[i] ?? Infinity,
      c = nums2[j - 1] ?? -Infinity,
      d = nums2[j] ?? Infinity;
    // 4. The cuts are correct when no left value is above a right value
    if (a <= d && c <= b) {
      // 5. Read the median off the largest left value and the smallest right one
      return (m + n) % 2
        ? Math.max(a, c)
        : (Math.max(a, c) + Math.min(b, d)) / 2;
    }
    // 6. Left value of nums1 above right value of nums2: take fewer from nums1
    if (a > d) {
      r = i - 1;
    } else {
      // 7. Otherwise the left value of nums2 is too big: take more from nums1
      l = i + 1;
    }
  }
  // Sorted arrays always have correct cuts, so this line is never reached.
  return 0;
}

def solve(nums1, nums2):
    # 1. Search the shorter array for how many values it gives the left half
    if len(nums1) > len(nums2):
        return solve(nums2, nums1)
    m, n = len(nums1), len(nums2)
    low, high = 0, m
    while low <= high:
        # 2. Try the middle count; nums2 gives the rest of the left half
        i = (low + high) // 2
        j = (m + n + 1) // 2 - i
        # 3. Note the values next to each array's cut, or infinity past an end
        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")
        # 4. The cuts are correct when no left value is above a right value
        if left1 <= right2 and left2 <= right1:
            # 5. Read the median off the largest left value and the smallest right one
            if (m + n) % 2:
                return max(left1, left2)
            return (max(left1, left2) + min(right1, right2)) / 2
        # 6. Left value of nums1 above right value of nums2: take fewer from nums1
        if left1 > right2:
            high = i - 1
        else:
            # 7. Otherwise the left value of nums2 is too big: take more from nums1
            low = i + 1
    # Sorted arrays always have correct cuts, so this line is never reached.
    return 0

Example 2, traced

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

Output
2.5

Why it works

As nums1 gives more values, its left value can only rise and the right value of nums2 can only fall. So the counts form a run of "too few", then the correct ones, then a run of "too many", and each test shows on which side the correct counts are. With correct cuts the left half holds exactly the smallest values, and the median is read at its edge.

Time and space

Time O(log min(m,n)), extra space O(1). The search covers the counts from 0 to the shorter array's length, min(m,n), and drops at least half of them each round. That is about log min(m,n) rounds of at most four reads each: O(log min(m,n)) time, which meets the task's aim. Only a few numbers are stored: O(1) space.

n and m are input sizes, h is a tree's height, V and E are a graph's vertices and edges. The bounds describe this reference solution.

Watch out

• The left half takes half of the total, rounded up. If you round down, example 1 returns 1 instead of 2: the left half then holds one value of three, and the median ends up in the right half.
• Search the shorter array. If you search the longer one, [-2, 0, 1, 2, 4] with an empty second array asks that empty array for a value: JavaScript returns 0 instead of 1, and Python stops with an error.
• Compare with "at most", not "below", because equal values may sit on both sides of a cut. With "below", [-1, 1] and [-3, 1, 1] has no passing cuts: the loop ends and the code returns 0 instead of 1.

Pattern

Binary search on the answer. Builds on Binary Search and Koko Eating Bananas.

Keep going