Binary search with confidence

Consider the following binary search code, copied from Wikipedia. Is it correct?

function binary_search(A, n, T):
  L := 0
  R := n − 1
  while L <= R:
    m := floor((L + R) / 2)
    if A[m] < T:
      L := m + 1
    else if A[m] > T:
      R := m - 1
    else:
      return m
  return unsuccessful

How sure are you?

When I see this code, I have to ask myself:

1. What’s up with all of these ‘s?
2. What happens on small arrays?
3. What happens if the element isn’t in the array?

Checking all of those cases can be tedious. I find a different implementation of binary search easier to write, understand, and prove correct. And once you understand the approach, you’ll also be able to implement binary search more confidently.

An outline

Loop invariants

Before we dive into the binary search algorithm, I want to review loop invariants. Later, we will use a loop invariant to help prove that our binary search implementation is correct.

Informally, loop invariants are properties of a loop that remain unchanged as the loop executes. For example, consider the following loop:

int f() {
  int num = 0;
  for (int i = 0; i < 100; ++i) {
    num += 2;
  }
  return num;
}

One invariant that remains true during the loop is that num is even (). Here’s the argument:

1. : num is initially even (since 0 is even).
2. : If num is even at the beginning of a loop iteration, the loop will add two to it; hence num will be even at the end of that loop iteration.

From these two statements, we can deduce that num must be even after the th loop iteration for any arbitrarily large : simply start from Statement 1, and repeatedly apply times:

1. (After loop iterations,) num is even. ()
2. After loop iteration, num is even. (Apply to 1)
3. After loop iterations, num is even. (Apply to 2)
4. After loop iterations, num is even. (Apply to 3)
5. …
6. After loop iterations, num is even.
7. After loop iterations, num is even. (Apply to 6)
8. Ad infinitum…

The two combined complete our argument and prove holds over all loop interations. In other words, is a loop invariant.

The above argument is an argument by mathematical induction. To make an argument by mathematical induction, one has to prove two things: the basis step (sometimes called the base case) and the inductive step. To prove that always held, was the basis step, and was the inductive step.

Both the basis step and the inductive step are important! Otherwise, the argument is not valid. For example, if num was initially , then the statement in our basis step could no longer be “num is initially as even.” We might instead choose an alternative statement, such as “num is initially odd.” That would change our invariant: num would always be odd, not even since a odd number would stay odd after adding two on each iteration. Similarly, if we had added a different constant to num every loop iteration, that would change the statement of our inductive step. Depending on what different constant we chose, the parity of num might not stay the same throughout the loop.

Reframing binary search

Returning back to binary search, suppose we want to find the index of 6 in the following sorted array:

Let’s consider the same array, except we label each element with a color—red or green. Color elements with value strictly smaller than 6 green, and elements with value greater than or equal to 6 red.

Notice how the array starts off as green (elements ), and then transitions to red (elements ). Also, notice that 6 is the first red element. The previous element, 3, is a green element. This informs our approach: instead of looking for 6 directly, we could instead look for the boundary between the green and red regions. Once we find the boundary, we can use it to locate the element we are searching for.

In addition, we can generalize this observation to all sorted arrays because they all share this structure. For any sorted array, the leftmost elements are colored green because they are less than 6. Then, traversing left to right, at some point some element will be 6 or larger. Starting from that point, the rest of the elements will be colored red. (This structure—turning red and staying red—does not appear in unsorted arrays; that’s why binary search requires a sorted array!)

The algorithm

We start with two pointers into the array, left and right. left points to the first element and right points to the last element:

This motivates our loop invariant, which is: left always points to a green element, and right always points to a red element. If we can keep moving the left and right pointers closer to each other while maintaining the invariant, then we will eventually find the boundary.

Let’s consider the middle element:

We have to choose a pointer—either left or right—to move to mid. Which pointer should we choose?

To help us decide, let’s look at our loop invariant. Since the mid element is a red element, we want to move right to mid. That would maintain our invariant that right always points to red. (Moving left to mid would break our invariant: our invariant says that left always points to green, but moving left to mid would make it point to red!)

Here are the pointers after we move left ← middle:

We pick the new middle element halfway between left and right:

Since this new middle element is green, and our invariant says that left points to green, we know we must move left to mid. Notice how our reasoning is purely mechanical: just look at the color of the middle element, and move pointer that must point to that color.

We repeat this process until left and right are adjacent. Here’s the next mid:

And it’s red, so we move right:

Now left and right are right next to each other. Since right is red, we know that it points to the first element greater than or equal to 6. Finally, we return the index that right points to (3).

The code

Our binary search function will take an array and a function is_green, which tells us whether an element is green or not. (If an element isn’t green, it must be red.)

is_green(0)  # True, because array[0] == 0  and 0 < 6
is_green(2)  # True, because array[2] == 3  and 3 < 6
is_green(3)  # False,  because array[3] == 6  and 6 >= 6
is_green(7)  # False,  because array[7] == 58 and 58 >= 6

Here is our binary search algorithm, implemented in Python.

def binary_search(array, is_green):
  left, right = 0, len(array) - 1

  # Main loop which narrows our search range.
  # All we do is check whether the middle element is green or not.
  # - If it's green, we move the left pointer.
  # - If it's not green, it's red, so we move the right pointer.
  while right - left > 1:
    mid = (left + right) // 2
    if is_green(array[mid]):
      left = mid
    else:
      right = mid

  return right

# Calling binary search on our array to look for 6.
binary_search(array, lambda x: x < 6);

Does this work? Recall our two desired invariants:

1. left points to a green element.
2. right points to a red element.

The above loop body preserves those invariants clearly: we only assign green ← mid if mid is green; otherwise, it is not green, so we assign red ← mid. Just by remembering these two invariants, you can easily reconstruct the loop body. There is no equality check, no +1/-1 arithmetic: just move the correct-color pointer.

The while condition computes the number of elements remaining in the search range (right - left) and keeps looping as long there is more than one element in that range [source]. Combined with our invariants, that means if the loop terminates, then left will point to the last green element and right will point to the first red element. When we exit the loop, we return the index of the first red element, which we know right stores.

The base case strikes back

Are we done? Not so fast!

In the previous section we showed that the loop body maintains our desired invariants. In other words, we proved the inductive step: that if left (right) pointed to a green (red) element at the beginning of a loop iteration, it would continue to point to a green (red) element at the end.

But that’s a big if, since we never established that the invariants were true in the first place! If the entire array was green, for instance, we would only move left. Then right would point to a green element after the loop ended.

Stepping back, to prove that our loop is correct, we must prove that our desired statements always hold. We’ve already proved the inductive step above. But the inductive step says that if our invariant held, then it continues to hold. It doesn’t tell us anything about whether our invariant was true in the first case! That is, to complete our argument, we must also show that our invariant was initially true—we must prove the basis step als holds.

To prove that basis step holds, we must ensure that before we enter the loop that our invariants are indeed satisfied. So we check that left points to green and right points to red. If either isn’t true, we can return early.

def binary_search(array, is_green):
  left, right = 0, len(array) - 1
  if not array:
    return ?
  if not is_green(array[left]):
    return ?
  if is_green(array[right]):
    return ?

  # Loop...

But what values should we return? There are three cases. First, if the array is all red, it’s clear: the first red element would be at index 0, so we can directly return 0.

Second, if the array is all green, then what we should return is not as clear. In practice, it’s useful to return an index one past the end of the array. One could imagine that the first red element (represented here as ) would be there:

Finally, if the array is empty, we choose to return 0 for reasons that will become clear soon.

Here’s the final code:

def binary_search(array, is_green):
  left, right = 0, len(array) - 1
  if not array:
    return 0
  if not is_green(array[left]):
    return 0
  if is_green(array[right]):
    return len(array)

  while right - left > 1:
    mid = (left + right) // 2
    if is_green(array[mid]):
      left = mid
    else:
      right = mid

  return right

# Calling binary search on our array to look for 6.
binary_search(array, lambda x: x < 6);

What happens if the target element is missing?

Consider the same array, except with 6 changed to 7. Let’s repeat our algorithm, searching for 6. I’ll omit the steps with “middle” for brevity.

We end up with the same return value. This is because our algorithm doesn’t actually find the index of 6—it finds the leftmost index at which you could insert 6 to keep the array sorted. That is, we could run the pseudocode:

# True.
assert is_sorted(array)

index = binary_search(array, lambda i: array[i] < 6)
array.insert(index, 6)

# Always still True.
assert is_sorted(array)

This is a feature—in some cases we only need to insert into a sorted array, so we don’t need to check if the returned index has a certain element in the binary search procedure. We can leave that to the caller.

This is also why we return one past the end if the array is entirely green, or 0 (which is also one past the “end” of the array) when the array is empty. In both cases, the returned index is the correct place to insert the new element.

Wikipedia’s binary search

Using invariants, can we prove the original Wikipedia code correct? Yes! Except the invariant is a bit more complicated. Here’s the original code again, where A is the array, n is the length of the array, and T is the desired element.

function binary_search(A, n, T):
  L := 0
  R := n − 1
  while L <= R:
    m := floor((L + R) / 2)
    if A[m] < T:
      L := m + 1
    else if A[m] > T:
      R := m - 1
    else:
      return m
  return unsuccessful

An invariant that works here is: “if T is in the array, then it must be at some index where .” Using that invariant, one can prove that if the algorithm returns an index m, A[m] = T; otherwise, T is not in the array and the algorithm returns unsuccessful. But not only is the proof more complicated than our proof above, its also left to the reader.