Median of 2 sorted arrays of equal size

The problem, simply

You get two sorted arrays A and B, and both have the same length n. If you poured them into one sorted array, it would have 2n numbers. That is always an even count, so the median is the average of the two middle numbers: the n-th and the (n+1)-th smallest.

Give back that median, faster than pouring them together.

A story to remember it

Picture it

Examples

Way 1: the simple way

Pour both rows together with two fingers, then read positions n - 1 and n.

1. Put one finger at the start of each row.
2. Copy the smaller number under your fingers into merged and move that finger.
3. When one row ends, copy the rest of the other.
4. Return the average of merged[n - 1] and merged[n].

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

How long it takes: every number is copied once, O(n). Extra room: the merged row, O(n). You could stop copying once you reach position n and keep only the last two numbers, which needs O(1) room but still O(n) time.

Way 2: the clever way

Choose how many numbers A puts on the short side, called cut1. Then B puts cut2 = n - cut1.

1. Search cut1 between lo = 0 and hi = n. Try the middle.
2. Read the four border numbers: left1 = A[cut1 - 1], right1 = A[cut1], left2 = B[cut2 - 1], right2 = B[cut2]. Past an end, use minus infinity on the left and plus infinity on the right.
3. If left1 <= right2 and left2 <= right1, the split is right. The median is (max(left1, left2) + min(right1, right2)) / 2.
4. If left1 > right2, A sent too many: hi = cut1 - 1. Otherwise A sent too few: lo = cut1 + 1.

Blue boxes are on the short side. Dashed boxes are the borders being checked, and a red box is the border that broke the rule.

class Solution {
    public double findMedian(int[] A, int[] B) {
        int n = A.length;
        int lo = 0, hi = n;
        while (lo <= hi) {
            int cut1 = (lo + hi) / 2;     // how many A sends to the short side
            int cut2 = n - cut1;          // B sends the rest
            int left1  = (cut1 == 0) ? Integer.MIN_VALUE : A[cut1 - 1];
            int right1 = (cut1 == n) ? Integer.MAX_VALUE : A[cut1];
            int left2  = (cut2 == 0) ? Integer.MIN_VALUE : B[cut2 - 1];
            int right2 = (cut2 == n) ? Integer.MAX_VALUE : B[cut2];
            if (left1 <= right2 && left2 <= right1) {
                return (Math.max(left1, left2) + (double) Math.min(right1, right2)) / 2.0;
            }
            if (left1 > right2) hi = cut1 - 1;   // A sent too many
            else lo = cut1 + 1;                  // A sent too few
        }
        throw new IllegalArgumentException("arrays were not sorted");
    }
}

How long it takes: each try halves the choices for cut1, O(log n). Extra room: a few numbers, O(1).

Another way: compare the two middles

This is the classic answer for the equal size version. Find the median of each row, m1 and m2.

• If they are equal, that is the answer.
• If m1 < m2, the true median is not in the small half of A and not in the big half of B. Throw both away (the same count from each, so the rows stay equal) and repeat.
• When each row has 1 or 2 numbers left, finish by hand.

class Solution {
    public double findMedian(int[] A, int[] B) {
        return solve(A, 0, A.length - 1, B, 0, B.length - 1);
    }

    private double middle(int[] a, int s, int e) {
        int len = e - s + 1;
        if (len % 2 == 1) return a[s + len / 2];
        return (a[s + len / 2 - 1] + (double) a[s + len / 2]) / 2.0;
    }

    private double solve(int[] A, int s1, int e1, int[] B, int s2, int e2) {
        int n = e1 - s1 + 1;
        if (n == 1) return (A[s1] + (double) B[s2]) / 2.0;
        if (n == 2) return (Math.max(A[s1], B[s2]) + (double) Math.min(A[e1], B[e2])) / 2.0;
        double m1 = middle(A, s1, e1), m2 = middle(B, s2, e2);
        if (m1 == m2) return m1;
        int drop = (n - 1) / 2;                  // same count from both rows
        if (m1 < m2) return solve(A, s1 + drop, e1, B, s2, e2 - drop);
        return solve(A, s1, e1 - drop, B, s2 + drop, e2);
    }
}

This is also O(log n) time, and O(log n) room for the waiting calls. All the math is done in double. An integer version of this code returns 3 instead of 3.5 for [1, 2, 3] and [4, 5, 6], a bug that hides in many copies online.

Watch out for

• The total 2n is always even, so there is no odd case here. In the different size problem there is.
• Past the end of a row, use plus or minus infinity. Without them you read A[-1] or A[n].
• Ties are fine. The checks use <=, so equal numbers on both sides of a cut do not break anything.

What interviewers ask next

• What if the sizes differ? Search the cut in the shorter row and let the other cut fill up to half. The same border check works.
• Why throw away the same count from both rows in the middles method? So the median of what is left stays the median of the whole: we drop as many numbers below it as above it.
• Find the k-th smallest of the two rows. Use k instead of n for the short side size, and clamp the search so cut2 stays inside B.
• Can it be done in O(1)? No. In the worst case you must look at about log n numbers to know where the middle is.

Remember this

Grown-up words: the short side and tall side are the two halves of a partition, the borders are left1, right1, left2, right2, and comparing the two middles and dropping halves is divide and conquer.