Binary search done right: finding the median of a union in logarithmic time

Previous post: Binary search done right: introduction

This is a LeetCode problem #4 (”hard”)

💡

In this problem, you are given two sorted arrays of lengths n and m, and you need to find the median of their union as fast as possible. Recall that if an array has odd length [1,2,3], then the median is the element 3 situated in the middle, while for arrays of even lengths [1,2,3,4], the median is a half-sum of the two middle elements (2+3)/2=2.5 .

The required algorithm has a complexity O(log(n+m))

Exercise. Can you prove that is the lower bound for complexity?

Solution to the exercise (click to expand).

This is proven in a similar way as the lower bound is proven for sorting algorithm. In the latter case, the output is one of possibilities which needs at least bits information. In the current scenario, the median may be any of the elements: strictly speaking, not any of the positions could indicate the median of the union, but almost any. Therefore, we need at least bits of information in asymptotics.

There are two parts in this problem: first, come up with a solution, and the second, code it without messing around with +/- 1 . The second part may turn out to be surprisingly tricky, but with proper approach it can be handled without too much difficulties.

Idea of the solution using binary search

In the example above, the total size of the union is , and so the half length is 19//2=9 . Their union is an array displayed below with the median marked.

[1, 1, 2, 3, 4, 5, 5, 6, 6, 7, 8, 11, 12, 15, 17, 18, 19, 20, 21]
                            ^

In fact, this selection of the first 9 elements corresponds to a union of the prefixes of these two arrays, respectively of lengths 3 and 6.

Note that the union of two prefixes does not necessarily correspond to a prefix of the union. On the contrary, any prefix of the union is indeed a union of two prefixes.

Convention. There are two cases for the total length, the odd one and the even one. Instead of coding a separate algorithm for these two cases, let us assume that the prefix we are searching for has a length (n+m)//2. In both cases, we will need to find the next element to the prefix, and only then, depending on whether the total length is even or odd, either return this next element, or return the half-sum of this element with the last element of the prefix.

Since the length of the median is fixed, we need to find the length of the second prefix only. The length of the first prefix is calculated as the half-length minus the length of the second prefix.

The solution is to perform the binary search on the length of the second array. How do we identify that we have obtained a good length? Well, the correct partitioning should identify a following condition: the last element of the second prefix is smaller than the next element after the first prefix.

If the length is too large, this condition is still satisfied. If the length is too small, this condition is broken. This is a perfect setting for the binary search where we identify the minimal length satisfying this condition (or, alternatively, the maximal length, not satisfying this condition).

We formulate a property, and at the end of the binary search, we have the left and the right ends, for which this property is, respectively, unsatisfied and satisfied.

Implementation

Here are several implementation aspects to emphasise.

The next element after the prefix union may be situated in the first or in the second array. This element is chosen by selecting the minimum among the two.

The same concern goes for the last element of the prefix union. The last element is obtained by taking the maximum.

In some cases (for example, when the two arrays have equal lengths), the length takes extreme/corner cases such as 0 or the total length of the second array. In some of the cases, the “next” element in the first array is unidentified because the list index is out of bounds.

Calculation of the position of the element in the first array requires some care.

The third point is the most important as it leads to the most bugs in implementation. The natural way to deal with sorted arrays is to define the auxiliary elements at positions -1 and len(array) to be -infty and +infty . These quantities are represented by float datatype. In order to access these elements, we can implement a separate function to avoid code duplication.

def access_elem(array, idx):
    if idx < 0:
        return -float("inf")
    if idx >= len(array):
        return float("inf")
    return array[idx]

Now, before going to the implementation, let us make the preliminary calculations.

The half-length is going to be .

If the length of the second prefix is , then the length of the first prefix is .

The position of the last element in the second prefix is .

Remark. We are subtracting a one here. This may seem not very nice, but if we keep in mind this is length-to-position conversion, this makes more sense.

The position of the next element after the first prefix is .

Remark. Despite innocent-looking subtraction of the two numbers, this is still length-to-position conversion, although this time the position is taken to be after the interval.

Without loss of generality, assume that the second array is not longer than the second.

We don’t use variable names n1 and n2 because they lead to more bugs. We use first and second instead, even though they need more characters.

Finally, recall the property for the binary search that is satisfied for the right half:

Property. The head of the second prefix is greater than the next element after the first prefix.

Final preparation: initial values. We need to guarantee that for the left and right boundaries, the property will or will not be satisfied. One can remark that when left = 0 we have the desired inequality since for any . However, len(second) might not be a sufficient upper bound for the length and we need to take len(second)+1 instead.

def findMedianSortedArrays(self, first: List[int], second: List[int]) -> float:
    if len(first) < len(second):
        first, second = second, first

    half = (len(first) + len(second)) // 2

    len_left = 0
    len_right = len(second)
    while len_right - len_left > 1:
        len_mid = (len_left + len_right) // 2

        pos_second = len_mid - 1
        pos_first = half - len_mid

        if access_elem(second, pos_second) > access_elem(first, pos_first):
            len_right = len_mid
        else:
            len_left = len_mid

		...

After finishing this binary search, we have len_left which is the target length for which we have second[pos_second] <= first[pos_first] . Now we need to calculate the four elements

the heads of the two prefixes

the next elements after each of the prefixes

and depending on the parity of the total length, output the min between the next elements, or the half-sum

Even if making the code explicit means more code, we prefer the latter since it is more reliable and requires less thinking.

...

# continuation
# assume we know `len_left` and `len_right`
# for simplicity of implementation, we ignore `len_right`

len_second = len_left
len_first = half - len_left

second_head = access_elem(second, len_second - 1)
second_next = access_elem(second, len_second)

first_head = access_elem(first, len_first - 1)
first_next = access_elem(first, len_first)

# here, we calculate the ultimate head and next
# we can't use the name "next" because this keyword is reserved

head_elem = max(first_head, second_head)
next_elem = min(first_next, second_next)

total_length = len(first) + len(second)

if total_length % 2 == 0:
    return (head_elem + next_elem) / 2
else:
    return next_elem

This code got written without any testing and got accepted after the first submit.

Aftermath

Even though the code is a bit lengthy, the critical part (the binary search) is kept really simple and solid.

The long part after the binary search is simply doing arithmetic manipulations between four elements. All the extra variables are introduced for human convenience.

The parts where we keep the left and the right elements out-of-borders might look counter-intuitive, but remember that we never access the left and the right

Next post: Binary search done right: longest increasing subsequence in O(n log n)