Median of Two Sorted Arrays (LeetCode 4)

The problem

You get two sorted arrays a (length m) and b (length n). Return the median of all m + n numbers combined. For an odd total, that’s the middle value; for an even total, it’s the average of the two middle values. The required time is O(log(m + n)).

Input:  a = [1, 2], b = [3, 4]
Output: 2.5

Combined, the numbers are [1, 2, 3, 4]. The two middle values are 2 and 3, so the median is 2.5.

How to think about it

Start with brute force. Merge the two arrays like merge sort and read the middle: O(m + n). Stopping the merge halfway saves space but not time. The problem wants logarithmic time.

Think about what the median does. It splits the combined numbers into a left half and a right half of equal size (the left gets one extra for an odd total), with everything on the left <= everything on the right. Describe that split without merging: the left half is the first i elements of a plus the first j of b, with i + j = (m + n + 1) // 2. Pick i and j is fixed.

When is a partition correct? Each array is sorted on its own, so the only checks that can fail are across the arrays. With al = a[i-1], ar = a[i], bl = b[j-1] and br = b[j], the partition is valid when al <= br and bl <= ar. Then the median is max(al, bl) for an odd total, or the average of max(al, bl) and min(ar, br) for an even total. Missing neighbours at the ends count as minus or plus infinity.

The key insight: binary-search i. If al > br, you took too many from a, so move i left (hi = i - 1). If bl > ar, you took too few, so move i right (lo = i + 1). That’s a monotonic decision, so binary search over i in [0, m] works. Search over the smaller array so j never goes out of range and the cost is O(log min(m, n)).

Which pattern is this? Binary search over a partition point rather than a value. The tell is “two sorted arrays” plus a log-time requirement for an order statistic (median, k-th smallest).

Approach

1. Swap so that a is the shorter array. Set lo = 0, hi = m.
2. While lo <= hi:
   • i = (lo + hi) // 2, j = (m + n + 1) // 2 - i.
   • Read al, ar, bl, br, using -inf / +inf when an index is off the end.
   • If al <= br and bl <= ar: return max(al, bl) if the total is odd, else (max(al, bl) + min(ar, br)) / 2.
   • Else if al > br: hi = i - 1. Else: lo = i + 1.

Solution

def find_median_sorted(a, b):
    if len(a) > len(b): a, b = b, a                 # search the shorter array
    m, n = len(a), len(b); lo, hi = 0, m
    while lo <= hi:
        i = (lo + hi) // 2; j = (m + n + 1) // 2 - i  # i from a, j from b on the left
        al = a[i-1] if i > 0 else float('-inf')
        ar = a[i]   if i < m else float('inf')
        bl = b[j-1] if j > 0 else float('-inf')
        br = b[j]   if j < n else float('inf')
        if al <= br and bl <= ar:                   # valid partition found
            if (m + n) % 2: return max(al, bl)
            return (max(al, bl) + min(ar, br)) / 2
        elif al > br: hi = i - 1                    # too many from a
        else: lo = i + 1                            # too few from a
# find_median_sorted([1,3],[2]) -> 2
# find_median_sorted([1,2],[3,4]) -> 2.5

Dry run

On a = [1, 2], b = [3, 4] (no swap needed, m = n = 2, left half size 2):

The left half is [1, 2], the right half is [3, 4]. The total is even, so the answer is (max(2, -inf) + min(inf, 3)) / 2 = 2.5.

Complexity

• Time: O(log min(m, n)). The binary search runs over i in [0, m], where m is the shorter length.
• Space: O(1).

Edge cases and pitfalls

• One array empty: a becomes the empty one, i = 0, and the median comes straight from b.
• Not searching the shorter array can make j negative or larger than n, which breaks the index checks.
• The + 1 in (m + n + 1) // 2 puts the extra element on the left for odd totals, which is why the odd case returns max(al, bl).
• Infinity sentinels handle partitions at either end of an array. Without them you get index errors at i = 0 or i = m.

Related problems

• Merge Two Sorted Lists (LeetCode 21): the linear merge that the brute force uses.
• Find Median from Data Stream (LeetCode 295): the median of a stream, kept with two heaps.
• Kth Largest Element in an Array (LeetCode 215): another order-statistic problem, on one unsorted array.