Finding The Median Of Two Sorted Arrays

The reason I’m writing about this problem, is not because of this specific problem itself, but because it shows that sometimes it’s simpler to solve a more general problem, where the original problem is a private case, and then applying the solution to this private case.

Maybe this is counterintuitive at first: How can happen that solving the general case is simpler? As engineers, we’re used to think in the terms of trade-offs, either you solve a general problem (win) with great effort (lose), or you solve a private and specific case (lose) with little effort (win). It doesn’t feel right to win twice: Both solve the more general problem, and do it with less effort.

This intuition seems to not always be correct, as in the case we’re discussing here. Moreover, many times, understanding a more general problem gives a better insight into the nature of the problem than a private case.

The message I’m trying to pass in this article, is: If you’re facing a problem, and something in the problem feels too constrained, don’t hesitate to think of a more general problem, by relaxing some constraints and unbinding parameters.

Let’s dive into our specific problem:

The problem:

Given two sorted arrays, “a” and “b”, with sizes “sa” and “sb” respectively, find the median of the union of these arrays.

Complexity requirements: O(log(sa + sb)) in the worst case, both time and space.

You can try to solve the above problem at once. I tried, and it was painful. It doesn’t feel natural: Something in the problem (and in complexity requirements if you want to cheat) tells us that the solution should be recursive, but median doesn’t really fit with recursion (On what variable do we want to do the recursion?).
You can try to force it, I’m not sure that you won’t suffer.

Surprisingly (or not?), the simplest solution (among those I happened to see), is the solution of a more general problem:

The Select Problem

Given two sorted arrays, “a” and “b”, with sizes “sa” and “sb” respectively, find the k-th smallest element of the union of these arrays. Under the same complexity requirements as above.

Once the select problem is solved, the median problem is solved as well, since the median is the ((sa + sb)/2 + 1)th smallest element (if sa+sb is odd), or the average of the two middle elements (if sa+sb is even).

So, let’s solve the select problem first:

The idea is to (virtually) partition “a” into two arrays: “a_left” and “a_right”, and the same for “b”: “b_left” and “b_right”, such that the size of “a_left” plus the size of “b_left” is exactly k.

Ideally, a_left will be a[0..(k/2-1)] and b_left will be b[0..(k/2-1)]. Unless, of course, a is smaller than k/2, in that case we’ll take “sa” elements from “a” (i.e. the whole array), and the rest of the elements (k – sa elements) from “b”.

Let’s assume that a_left has “naelemetns, and b_left has “nb” elements. As we stated before, na+nb = k.

Trivial case 1:

k = 1. i.e., we need to find the smallest element of the union of “a” and “b”. Obviously, it’s either the first element of “a” or the first element of “b”. This operation is cheap. O(1)

Trivial case 2:

a[na-1] = b[nb-1]  (we assume zero based indexing, that is, first element of a is a[0])

In this case, the k-th smallest element is a[na-1] (or b[nb-1]): We have exactly (na + nb = k) elements smaller than it.

Case 3:

a[na-1] < b[nb-1]

This is the interesting case. If a[na-1] is smaller than b[nb-1], then:

1. The k-th smallest element cannot be in the “a_left” array:

All elements in “a_right” are greater than any of those in “a_left” (“a” is sorted). Total: sana elements

All elements in “b_right” are greater than any of those in “a_left” (“b” is sorted and a[na-1] < b[nb-1]). Total: sbnb elements

At least one element in “b_left” is greater than any of those in “a_left” (we already know that b[nb-1] is such an element). Total: 1 (at least)

So, for any element in “a_left”, there are at least (sana + sbnb + 1 = sa + sb – k + 1) elements which are greater than it.

The k-th smallest element must have exactly (sa + sb – k) elements greater than it. Hence, none of the elements in “a_left” can be the k-th smallest element.

2. All of the elements in “a_left” are among the k smallest elements (but none is the k-th smallest) – Using the same logics above.

3. None of the elements in “b_right” are the k-th smallest element. We already know that all the elements in “a_left” and all the elements in “b_left” (total: k elements) are smaller than any element in “b_right”

From the three observation above, we can state that the k-th smallest element of the union of “a” and “b”, is the (k – na)th smallest element of the union of “a_right” and “b_left” (smells like recursion): We excluded “a_left” in (1), excluded “b_right” in (3), and we already have “na” of the k smallest elements, so we need to find the (k-na)th smallest in the arrays not excluded.

The last case, where a[na-1] > b[nb-1] is similar to the earlier case, except that “a” and “b” have exchanged their roles.

We covered all possible cases above. What is left is to translate it into some programming language, and then applying it to the private case where the k-th element is the median (remember?)

Here is the code:

Leave a Reply