Topics › Binary Search › BS on Answers

Median of 2 sorted arrays

Hard TUF ID 77 A2ZSDE

Problem

You are given two sorted arrays a (length n) and b (length m). Imagine merging them into one sorted array of n + m elements. Return the median of that merged array: the middle element when n + m is odd, or the average of the two middle elements when n + m is even. Return the result as a double.

Constraints: 0 ≤ n, m ≤ 1000, n + m ≥ 1, values fit in a 32-bit signed integer. The interesting requirement is to do it in O(log(min(n, m))) time.

Example 1

a = [1, 3], b = [2]
answer = 2.0         // merged = [1, 2, 3], the middle element is 2

Example 2

a = [1, 2], b = [3, 4]
answer = 2.5         // merged = [1, 2, 3, 4], the middle elements are 2 and 3

Example 3

a = [1, 4, 7, 10, 12], b = [2, 3, 6, 15]
answer = 6.0         // merged = [1, 2, 3, 4, 6, 7, 10, 12, 15], 9 elements, middle is 6

Approach 1: Brute force, merge and read the middle

Intuition. The problem statement literally describes an algorithm: merge the two sorted arrays with the standard two-pointer merge, then read the middle. Merging keeps both inputs’ order, so the result is sorted and the median sits at index (n + m) / 2 (and (n + m) / 2 − 1 for the even case).

Algorithm.

  1. Allocate merged of size n + m.
  2. While both pointers are in range, copy the smaller head and advance that pointer; then copy whatever remains.
  3. If the total is odd, return merged[total / 2]; otherwise return the average of merged[total / 2 − 1] and merged[total / 2].
static double bruteForce(int[] a, int[] b) {
    int n = a.length, m = b.length;
    int[] merged = new int[n + m];
    int i = 0, j = 0, k = 0;
    while (i < n && j < m) merged[k++] = a[i] <= b[j] ? a[i++] : b[j++];
    while (i < n) merged[k++] = a[i++];
    while (j < m) merged[k++] = b[j++];
    int total = n + m;
    if (total % 2 == 1) return merged[total / 2];
    return ((double) merged[total / 2 - 1] + merged[total / 2]) / 2.0;
}

Complexity. Time O(n + m) for the merge. Space O(n + m) for the merged copy.

Approach 2: Better, merge without storing

Intuition. We never need the merged array, only the one or two elements at the middle positions. Run the same merge, but instead of writing elements out, just count how many have been produced. When the count reaches a middle index, remember that element; after the second middle index there is nothing left to learn, so stop.

Algorithm.

  1. Let idx2 = total / 2 and idx1 = idx2 − 1 be the positions of interest in merged order.
  2. Walk the merge, producing one element per step and tracking its position pos.
  3. Save the element at idx1 into el1 and the one at idx2 into el2; break after idx2.
  4. Odd total: return el2. Even total: return (el1 + el2) / 2.
static double better(int[] a, int[] b) {
    int n = a.length, m = b.length, total = n + m;
    int idx2 = total / 2, idx1 = idx2 - 1;           // positions of the middle element(s) in merged order
    int el1 = 0, el2 = 0;
    int i = 0, j = 0, pos = 0;
    while (i < n || j < m) {
        int next;
        if (j >= m || (i < n && a[i] <= b[j])) next = a[i++]; else next = b[j++];
        if (pos == idx1) el1 = next;
        if (pos == idx2) { el2 = next; break; }
        pos++;
    }
    if (total % 2 == 1) return el2;
    return ((double) el1 + el2) / 2.0;
}

The condition j >= m || (i < n && a[i] <= b[j]) takes from a when b is exhausted, or when both have elements and a’s head is not larger. It never reads out of bounds because i < n is checked before a[i] and j >= m is checked before b[j].

Complexity. Time O(n + m): in the worst case the walk reaches the middle, about (n + m) / 2 steps. Space O(1).

This removes the memory cost but is still linear. To do better we have to stop walking and start cutting.

Approach 3: Optimal, binary search on the partition

Intuition. Think of the median differently: the merged array splits into a left half and a right half of fixed sizes, and the median is determined by the largest element of the left half and the smallest of the right half. A left half of size leftSize = (n + m + 1) / 2 must take some i elements from a (its first i, since a is sorted) and the remaining j = leftSize − i elements from b (its first j).

For a given i, the split is valid exactly when every element on the left is at most every element on the right. Since each array is sorted on its own, only two cross comparisons matter: the last of a’s left part (l1 = a[i-1]) must be at most the first of b’s right part (r2 = b[j]), and the last of b’s left part (l2 = b[j-1]) must be at most the first of a’s right part (r1 = a[i]).

If l1 > r2, we took too many from a: shift the cut left by lowering i. If l2 > r1, we took too few from a: raise i. Because both arrays are sorted, these conditions are monotone in i, so binary search over i in [0, n] finds the valid cut. Searching over the shorter array keeps j within [0, m] automatically and gives the O(log(min(n, m))) bound.

Algorithm.

  1. If a is longer than b, swap them, so the search runs over the shorter array.
  2. Let leftSize = (n + m + 1) / 2, lo = 0, hi = n.
  3. While lo <= hi: let i = mid, j = leftSize − i.
    • Compute l1, l2, r1, r2, using −∞ for an empty left part and +∞ for an empty right part.
    • If l1 <= r2 and l2 <= r1: odd total returns max(l1, l2); even total returns (max(l1, l2) + min(r1, r2)) / 2.
    • Else if l1 > r2: hi = i − 1. Otherwise lo = i + 1.
static double optimal(int[] a, int[] b) {
    if (a.length > b.length) return optimal(b, a);  // always search over the shorter array
    int n = a.length, m = b.length, total = n + m;
    int leftSize = (total + 1) / 2;                  // size of the left half; it gets the extra element when total is odd
    int lo = 0, hi = n;
    while (lo <= hi) {
        int i = lo + (hi - lo) / 2;                  // i elements of a go to the left half ...
        int j = leftSize - i;                        // ... so j elements of b must go there too
        int l1 = i == 0 ? Integer.MIN_VALUE : a[i - 1];
        int l2 = j == 0 ? Integer.MIN_VALUE : b[j - 1];
        int r1 = i == n ? Integer.MAX_VALUE : a[i];
        int r2 = j == m ? Integer.MAX_VALUE : b[j];
        if (l1 <= r2 && l2 <= r1) {
            if (total % 2 == 1) return Math.max(l1, l2);
            return ((double) Math.max(l1, l2) + Math.min(r1, r2)) / 2.0;
        }
        if (l1 > r2) hi = i - 1;                     // too many from a: move the cut left
        else lo = i + 1;                             // too few from a: move the cut right
    }
    throw new IllegalArgumentException("inputs are not sorted");
}

Dry run on a = [1, 4, 7, 10, 12], b = [2, 3, 6, 15]. a is longer, so the roles swap: search over a = [2, 3, 6, 15] (n = 4) against b = [1, 4, 7, 10, 12] (m = 5). total = 9, leftSize = 5.

lo hi i j l1 l2 r1 r2 valid? action
0 4 2 3 3 7 6 10 no: l2 = 7 > r1 = 6 lo = 3
3 4 3 2 6 4 15 7 yes odd total: max(6, 4) = 6

The median is 6.0, matching the merge.

Complexity. Time O(log(min(n, m))): binary search over the shorter array with constant work per probe. Space O(1).

Why does a valid cut always exist and why is it unique in effect? Take the real merged array and cut it after leftSize elements; that cut takes some prefix of a and some prefix of b, and it satisfies both cross conditions. So the search space contains a valid i. Different valid i can occur only with duplicate values, and then max(l1, l2) and min(r1, r2) are the same numbers regardless, so the returned median is the same.

Edge cases and pitfalls