Median of 2 sorted arrays of different size

The problem, simply

You get two sorted arrays, A with m numbers and B with n numbers. The sizes can be different, and one can even be empty. Imagine pouring them together into one sorted array. Give back its median: the middle number if the total is odd, or the average of the two middle numbers if the total is even.

The famous version asks you to do it in O(log(m + n)) time, so you may not really pour them together.

A story to remember it

Picture it

Examples

Way 1: the simple way

Pour both arrays into one with two fingers, like merging two sorted piles of cards. Then read the middle.

1. Put a finger at the start of A and of B.
2. Copy the smaller of the two numbers under your fingers into merged, and move that finger.
3. When one array runs out, copy the rest of the other.
4. Read the middle one or two numbers of merged.

class Solution {
    public double findMedianSortedArrays(int[] A, int[] B) {
        int m = A.length, n = B.length;
        int[] merged = new int[m + n];
        int i = 0, j = 0, k = 0;
        while (i < m && j < n) merged[k++] = (A[i] <= B[j]) ? A[i++] : B[j++];
        while (i < m) merged[k++] = A[i++];
        while (j < n) merged[k++] = B[j++];
        int total = m + n;
        if (total % 2 == 1) return merged[total / 2];
        return (merged[total / 2 - 1] + merged[total / 2]) / 2.0;
    }
}

How long it takes: every number is copied once, O(m + n). Extra room: the merged array, O(m + n). You can drop the array by only counting up to the middle and remembering the last two numbers, which makes the room O(1), but the time is still O(m + n).

Way 2: the clever way

Do not pour. Cut each array into a left part and a right part.

1. Make sure A is the shorter one (swap if needed). We will search over cuts in A only.
2. The left side must hold half = (m + n + 1) / 2 numbers. If A gives cutA numbers, then B must give cutB = half - cutA.
3. Look at the four numbers around the cuts: leftA, rightA, leftB, rightB. A missing number past an end counts as minus infinity on the left or plus infinity on the right.
4. The cut is right when leftA <= rightB and leftB <= rightA. Then every left number is at most every right number.
5. If leftA > rightB, A gave too many: move the cut in A to the left. Otherwise A gave too few: move it to the right.
6. When it is right: for an odd total, the median is max(leftA, leftB). For an even total, it is the average of max(leftA, leftB) and min(rightA, rightB).

Each guess halves the choices for cutA, which is binary search. Blue boxes below are on the left side of a cut. Dashed boxes are the four border numbers being checked, and a red box is the border that broke the rule.

class Solution {
    public double findMedianSortedArrays(int[] A, int[] B) {
        if (A.length > B.length) return findMedianSortedArrays(B, A); // search the shorter one
        int m = A.length, n = B.length;
        int half = (m + n + 1) / 2;          // how many numbers go on the left
        int lo = 0, hi = m;
        while (lo <= hi) {
            int cutA = (lo + hi) / 2;
            int cutB = half - cutA;
            int leftA  = (cutA == 0) ? Integer.MIN_VALUE : A[cutA - 1];
            int rightA = (cutA == m) ? Integer.MAX_VALUE : A[cutA];
            int leftB  = (cutB == 0) ? Integer.MIN_VALUE : B[cutB - 1];
            int rightB = (cutB == n) ? Integer.MAX_VALUE : B[cutB];
            if (leftA <= rightB && leftB <= rightA) {         // the cut is right
                if ((m + n) % 2 == 1) return Math.max(leftA, leftB);
                return (Math.max(leftA, leftB) + (double) Math.min(rightA, rightB)) / 2.0;
            } else if (leftA > rightB) {
                hi = cutA - 1;                                // A gave too many
            } else {
                lo = cutA + 1;                                // A gave too few
            }
        }
        throw new IllegalArgumentException("arrays were not sorted");
    }
}

How long it takes: binary search over the cuts of the shorter array, O(log(min(m, n))). Extra room: a few numbers, O(1).

Watch out for

• Use (m + n + 1) / 2 for the left size. Then for an odd total, the median sits on the left side, and max(leftA, leftB) is the answer.
• An empty array is fine: m = 0 means cutA is always 0 and the infinities do the rest.
• Adding two big int values can overflow. Cast one to double (or long) before adding, as in the code.
• Return a double. Integer division turns 2.5 into 2.

What interviewers ask next

• Find the k-th smallest of two sorted arrays. Same cut idea with half replaced by k, or compare A[k/2 - 1] with B[k/2 - 1] and throw away k/2 numbers each step, O(log k).
• Why is the answer O(log(min(m, n))) and not O(log(m + n))? We only search the shorter array, so it is even better than asked.
• What if the arrays are equal in size? The same code works. There is also an older trick that compares the two medians; see the equal size problem.
• Median of a stream of numbers? That is a different tool: two heaps, one for the smaller half and one for the bigger half.

Remember this

Grown-up words: the cut is a partition, the left side is the lower half of the merged order, minus and plus infinity are sentinel values, and guessing then halving is binary search on the partition index.