DSA Median of Two Sorted Arrays

The problem

The median of two sorted arrays is the middle value of their combined order. Each array arrives already sorted. The answer must not merge them.

An odd total count answers the single middle value. An even total count answers the average of the two middle values. The cases below are the ones this article solves.

cases = [
    ([1, 3], [2], 2),
    ([1, 2], [3, 4], 2.5),
    ([1, 3, 5], [2, 4, 6], 3.5),
    ([], [2, 4], 3.0),
    ([2, 7, 11], [3, 9], 7),
]
for a, b, expected in cases:
    print("a", a, "b", b, "median:", expected)

Read one row at a time. The fourth row has an empty first array, so the median comes from the second array alone. The last row interleaves the two arrays in the combined order.

Intuition: merge the arrays and pick the middle

The first idea merges both arrays into one sorted list. The median then sits at the middle index. The merge below walks both arrays with two pointers.

a = [1, 3, 5]
b = [2, 4]
merged = []
i = 0
j = 0
while i < len(a) and j < len(b):
    if a[i] <= b[j]:
        merged.append(a[i])
        i += 1
    else:
        merged.append(b[j])
        j += 1
merged.extend(a[i:])
merged.extend(b[j:])
print("merged:", merged)
print("median:", merged[len(merged) // 2])
print("values stored to answer one question:", len(merged))

The merge is correct. Its cost is the trap. It copies and stores every value just to read one or two of them.

Brute force: the merge, counted

The brute force is the merge written as a function. A counter records every value copied into the merged list. The median rule reads the middle of that list.

def median_brute(a, b):
    merged = []
    copies = 0
    i = 0
    j = 0
    while i < len(a) and j < len(b):
        if a[i] <= b[j]:
            merged.append(a[i])
            i += 1
        else:
            merged.append(b[j])
            j += 1
        copies += 1
    while i < len(a):
        merged.append(a[i])
        i += 1
        copies += 1
    while j < len(b):
        merged.append(b[j])
        j += 1
        copies += 1
    total = len(merged)
    if total % 2 == 1:
        return merged[total // 2], copies
    return (merged[total // 2 - 1] + merged[total // 2]) / 2, copies

cases = [([1, 3], [2]), ([1, 2], [3, 4]), ([1, 3, 5], [2, 4, 6]), ([], [2, 4]), ([2, 7, 11], [3, 9])]
for a, b in cases:
    median, copies = median_brute(a, b)
    print("a", a, "b", b, "median:", median, "copies:", copies)

The copy count equals the total number of values every time. Nothing is saved on easy inputs. A million values cost a million copies for one middle value.

The optimal solution: binary search over a partition

A partition is a cut that splits an array into a left part and a right part, and the right partition puts exactly half of all values on the left with every left value at most every right value. That one sentence is the whole method.

One cut in each array builds the combined left and right sides. Only the four boundary values need comparing. A cut that is too far right or left fails one comparison, and binary search moves it.

An empty left or right part reads a sentinel instead of a value. The sentinels are minus infinity on the left and plus infinity on the right, so edge cuts compare cleanly.

def find_median(a, b):
    if len(a) > len(b):
        a, b = b, a
    m = len(a)
    n = len(b)
    half = (m + n + 1) // 2
    lo = 0
    hi = m
    while lo <= hi:
        i = (lo + hi) // 2
        j = half - i
        a_left = a[i - 1] if i > 0 else float("-inf")
        a_right = a[i] if i < m else float("inf")
        b_left = b[j - 1] if j > 0 else float("-inf")
        b_right = b[j] if j < n else float("inf")
        if a_left <= b_right and b_left <= a_right:
            if (m + n) % 2 == 1:
                return max(a_left, b_left)
            return (max(a_left, b_left) + min(a_right, b_right)) / 2
        if a_left > b_right:
            hi = i - 1
        else:
            lo = i + 1

cases = [
    ([1, 3], [2], 2),
    ([1, 2], [3, 4], 2.5),
    ([1, 3, 5], [2, 4, 6], 3.5),
    ([], [2, 4], 3.0),
    ([2, 7, 11], [3, 9], 7),
]
for a, b, expected in cases:
    result = find_median(a, b)
    print("a", a, "b", b, "->", result, "expected:", expected, "match:", result == expected)

All five cases pass. The search swaps the arrays first when needed, so the binary search always runs over the smaller array.

Walkthrough: finding the median of [1, 3] and [2]

The total count is three, so the left side must hold two values. The search tries a cut in the first array, and the second cut is forced by the half size. Each step prints both cuts and the verdict.

def traced_median(a, b):
    if len(a) > len(b):
        a, b = b, a
    m = len(a)
    n = len(b)
    half = (m + n + 1) // 2
    lo = 0
    hi = m
    steps = 0
    while lo <= hi:
        i = (lo + hi) // 2
        j = half - i
        steps += 1
        a_left = a[i - 1] if i > 0 else float("-inf")
        a_right = a[i] if i < m else float("inf")
        b_left = b[j - 1] if j > 0 else float("-inf")
        b_right = b[j] if j < n else float("inf")
        if a_left <= b_right and b_left <= a_right:
            print("step", steps, "cut a:", i, "cut b:", j, "verdict: holds")
            return max(a_left, b_left)
        if a_left > b_right:
            print("step", steps, "cut a:", i, "cut b:", j, "verdict: move left")
            hi = i - 1
        else:
            print("step", steps, "cut a:", i, "cut b:", j, "verdict: move right")
            lo = i + 1

print("median:", traced_median([1, 3], [2]))

The search runs over the shorter array, [2], after the swap at the top. Step one cuts it at 0 and the longer array at 2. The left side then holds 1 and 3, and 3 is larger than the 2 waiting on the right, so the cut moves right.

Step two cuts at 1 and 1. The left side holds 2 and 1, and the right side holds 3 alone. The partition holds, and the median is the largest left value, 2.

Complexity

The partition search halves the smaller array on each step, so its time is O(log n) in the smaller length. It stores only the cut positions and boundary values, so its extra space is O(1).

The merge copies every value once, so its time and extra space are both O(n) in the total count. The counts below run both methods on growing arrays, with the whole first array below the second.

def median_counted(a, b):
    if len(a) > len(b):
        a, b = b, a
    m = len(a)
    n = len(b)
    half = (m + n + 1) // 2
    lo = 0
    hi = m
    steps = 0
    while lo <= hi:
        i = (lo + hi) // 2
        j = half - i
        steps += 1
        a_left = a[i - 1] if i > 0 else float("-inf")
        a_right = a[i] if i < m else float("inf")
        b_left = b[j - 1] if j > 0 else float("-inf")
        b_right = b[j] if j < n else float("inf")
        if a_left <= b_right and b_left <= a_right:
            return steps
        if a_left > b_right:
            hi = i - 1
        else:
            lo = i + 1

for size in [5, 50, 500]:
    a = list(range(size))
    b = list(range(size, 2 * size))
    steps = median_counted(a, b)
    print("values per array:", size, "- merge copies:", 2 * size, " partition steps:", steps)

At 500 values per array the merge pays 1,000 copies. The search pays a handful of steps, growing by roughly one step per doubling. That slow growth is the logarithm at work.

Edge cases and pitfalls

Empty arrays are the main edge case. The swap at the top and the sentinels handle them with no special branch. The runs below cover both empty positions and a single value overall.

def find_median(a, b):
    if len(a) > len(b):
        a, b = b, a
    m = len(a)
    n = len(b)
    half = (m + n + 1) // 2
    lo = 0
    hi = m
    while lo <= hi:
        i = (lo + hi) // 2
        j = half - i
        a_left = a[i - 1] if i > 0 else float("-inf")
        a_right = a[i] if i < m else float("inf")
        b_left = b[j - 1] if j > 0 else float("-inf")
        b_right = b[j] if j < n else float("inf")
        if a_left <= b_right and b_left <= a_right:
            if (m + n) % 2 == 1:
                return max(a_left, b_left)
            return (max(a_left, b_left) + min(a_right, b_right)) / 2
        if a_left > b_right:
            hi = i - 1
        else:
            lo = i + 1

print("empty first array:", find_median([], [2, 4]))
print("empty second array:", find_median([1, 5], []))
print("one value total:", find_median([7], []))
print("all of a below all of b:", find_median([1, 2], [10, 20]))
print("negative values:", find_median([-5, -1], [-3]))
print("equal values:", find_median([2, 2], [2, 2]))
print("left sentinel:", float("-inf"), " right sentinel:", float("inf"))

The all-below row settles with the cut at the very end of the first array. Its right sentinel is what lets that cut pass the comparisons. Reading a real element past the end there is the classic crash in this problem.