Iterative binary search templates in Python

Why this matters

Binary search is one of the most used algorithms in programming. Any time code looks up a value in a sorted list, binary search is a good candidate for the work. Interview questions use it too, often with a twist: the answer is not a value in the list but a position, a count, or a boundary between two kinds of values. Most wrong answers to these questions come from small mistakes in the loop. A variable that never moves, a boundary that goes one step too far, or a middle position that rounds the wrong way.

The good news is that the loop has only three shapes. This lesson gives you each shape as a template. You learn when to use each one, why each line is there, and how to spot the mistake that makes it loop forever. The first lesson in this module explains why halving the range is fast. This lesson is about writing the loop so that it is correct.

The idea: keep a range that must contain the answer

Start with a game. A friend thinks of a number from 1 to 100. You guess 50. The friend says "higher". Now the answer must be between 51 and 100, so you never need to think about 1 to 50 again. Each guess removes half of what is left.

Binary search does the same thing to a sorted list. A list is Python's numbered collection of values. Each value has an index, which is its position. Python counts indexes from 0. A list of 7 values has indexes 0 to 6. A sorted list has its values in order, from smallest to largest. This lesson uses lists in the places where some languages would use an array.

Keep two variables:

• lo is the smallest index that could still be the answer.
• hi is the largest index that could still be the answer.

Each step of the loop looks at the middle index, called mid. It uses the result to move lo or hi so the range gets smaller. The loop ends when the range is empty or holds only one index.

Keep one rule while you write the loop. The answer, if there is one, is always inside the range from lo to hi, with both ends included. A statement that stays true on every step of a loop is called an invariant. If the invariant is true when the loop starts, and every step keeps it true, then the loop cannot throw away the answer by accident. Each template below has its own invariant, and each one is written next to the code.

The middle index is written like this:

mid = lo + (hi - lo) // 2

It is not written as (lo + hi) // 2. Other languages need the longer form to avoid a number limit. Python has no such limit, so both forms give the same answer here, as Mistake 5 shows. The longer form is kept as a habit that also works in those languages.

The // operator is integer division. It divides and drops the fraction, so 7 // 2 is 3. The single slash / gives a decimal number instead, so 7 / 2 is 3.5. Mistake 3 shows what goes wrong when / is used by mistake.

For now, read hi - lo as "how far apart the two ends are", and mid as "the start plus half of that distance".

Template 1: find an exact value

The first question is the simplest one. Is the target in the list, and at which index? Take the list [2, 4, 4, 4, 7, 9, 12] and the target 7. The answer is index 4.

Here is the search done by hand. The first range is indexes 0 to 6.

1. lo is 0 and hi is 6, so mid is 3. The value at index 3 is 4. That is less than 7, so the answer must be to the right. Set lo to 4.
2. The range is now 4 to 6, so mid is 5. The value at index 5 is 9. That is more than 7, so the answer must be to the left. Set hi to 4.
3. The range is now 4 to 4, so mid is 4. The value at index 4 is 7. That is the target. Return 4.

Each step checks the middle value and then throws away the half that cannot hold the target. Here is the template in Python:

def find_exact(a, target):
    lo = 0
    hi = len(a) - 1
    while lo <= hi:
        mid = lo + (hi - lo) // 2
        if a[mid] == target:
            return mid
        if a[mid] < target:
            lo = mid + 1
        else:
            hi = mid - 1
    return -1

Read the template line by line. hi starts at len(a) - 1 because that is the last index. The while condition is lo <= hi. When lo is greater than hi, the range is empty and the target is not in the list. The function then returns -1, which is a common sign for "not found" because no valid index is negative. Python also accepts a negative index, but then it means a position counted from the end, so a[-1] is the last value. A caller must check for -1 before using the result as an index. Mistake 6 shows what happens when it does not.

The two moves inside the loop are lo = mid + 1 and hi = mid - 1. Both skip mid itself. That is safe because the code has just compared mid with the target and knows it is not the answer. If the code wrote lo = mid instead, the range could stop shrinking. Exercise 4 in the practice section shows that case.

Here is the template as a complete program. Press Run. The program needs no input. Each line in the output shows one call and its result.

The program passes the question to the other two functions as a function. A predicate is a function that answers yes or no. In Python it returns True or False, and any function with that shape can be a predicate. The sections below explain how each one is used. The output of the program is:

find_exact 7 -> 4
find_exact 5 -> -1
first_true a[i] >= 4 -> 1
first_true a[i] > 4 -> 4
last_true a[i] <= 4 -> 3
last_true a[i] <= 1 -> -1

The first two lines match the hand trace above. The value 5 is not in the list, so the search returns -1.

Template 2: first true

Many questions are not "is the value here". They ask "where does this kind of value start?" For example, where does the first value that is 4 or more appear? In the list [2, 4, 4, 4, 7, 9, 12], the answer is index 1.

A predicate is a yes-or-no question about an index. In Python, the predicate for this question is written as a lambda, which is a short unnamed function: lambda i: a[i] >= 4. It reads "given index i, is the value at i at least 4?"

Look at the answers for every index in the list:

index:      0      1     2     3     4     5     6
value:      2      4     4     4     7     9    12
a[i] >= 4?  no    yes   yes   yes   yes   yes   yes

The answers go no, no, ..., yes, yes, yes. A sequence like this is called false then true, or monotone. The job of the template is to find the first yes. Here is the template:

def first_true(n, pred):
    lo = 0
    hi = n
    while lo < hi:
        mid = lo + (hi - lo) // 2
        if pred(mid):
            hi = mid
        else:
            lo = mid + 1
    return lo

Compare this with Template 1. Four things are different.

1. hi starts at n, not n - 1. The answer might be past the last index. If every answer is no, the first yes is "just after the end", which is index n. The function returns n in that case.
2. The loop condition is lo < hi, not lo <= hi. When lo equals hi, the range holds one index, and that index is the answer. There is nothing left to check.
3. There is no early return inside the loop. A yes at mid does not finish the search, because an earlier yes might exist. So the code moves hi to mid and keeps going.
4. A no at mid means mid and everything before it is too early. So the code moves lo to mid + 1.

The invariant for this template is: the first yes, or n if there is none, is inside the range from lo to hi. The line hi = mid keeps this true because mid might be the first yes. The line lo = mid + 1 keeps it true because mid is a no and cannot be the first yes.

Here is a hand trace with n = 7 and the predicate a[i] >= 4.

1. The range is 0 to 7. mid is 3. Index 3 holds 4, which is a yes. Set hi to 3.
2. The range is 0 to 3. mid is 1. Index 1 holds 4, which is a yes. Set hi to 1.
3. The range is 0 to 1. mid is 0. Index 0 holds 2, which is a no. Set lo to 1.
4. The range is 1 to 1, so the loop ends. Return 1.

Lower bound and upper bound

Two common questions use this template. The lower bound of a value is the first index where the value is not less than the target. It uses the predicate a[i] >= target. The upper bound is the first index where the value is greater than the target. It uses a[i] > target.

The difference between the two is the number of copies of the target in the list. In [2, 4, 4, 4, 7, 9, 12], the lower bound of 4 is index 1, and the upper bound of 4 is index 4. So there are 4 minus 1, which is 3 copies of 4.

The standard library module bisect has both searches built in. bisect_left returns the lower bound and bisect_right returns the upper bound. Run this to see the same two numbers:

The output is:

lower bound of 4: 1
upper bound of 4: 4

Template 3: last true

Now the pattern is reversed. The answers go yes, yes, ..., no, no, no. Here is the question "which index is the last one whose value is 4 or less?" For [2, 4, 4, 4, 7, 9, 12], the answer is index 3.

index:      0      1     2     3     4     5     6
value:      2      4     4     4     7     9    12
a[i] <= 4?  yes   yes   yes   yes    no    no    no

This template has one trap that the first two do not have. The middle index must round up, not down. Here is the code:

def last_true(n, pred):
    lo = 0
    hi = n - 1
    while lo < hi:
        mid = lo + (hi - lo + 1) // 2
        if pred(mid):
            lo = mid
        else:
            hi = mid - 1
    if n > 0 and pred(lo):
        return lo
    return -1

Look at mid first. The formula lo + (hi - lo + 1) // 2 gives the upper middle when the range has an even count of indexes. For a range from 3 to 4, the old formula lo + (hi - lo) // 2 gives 3 + 0, so mid is 3. The new formula gives 3 + 1, so mid is 4. This matters because a yes at mid sets lo = mid. If mid were 3 and lo were already 3, the range would not change at all.

Here is a hand trace with n = 7 and the predicate a[i] <= 4.

1. The range is 0 to 6. mid is 0 + (6 - 0 + 1) // 2, which is 3. Index 3 holds 4, which is a yes. Set lo to 3.
2. The range is 3 to 6. mid is 3 + (6 - 3 + 1) // 2, which is 5. Index 5 holds 9, which is a no. Set hi to 4.
3. The range is 3 to 4. mid is 3 + (4 - 3 + 1) // 2, which is 4. Index 4 holds 7, which is a no. Set hi to 3.
4. The range is 3 to 3, so the loop ends. The value at index 3 is a yes, so return 3.

The check after the loop, pred(lo), is needed when no index is a yes. If the predicate is false for every index, lo never moves from 0. The check then sees that index 0 is a no and returns -1. The n > 0 part handles an empty list, where hi is -1 and the loop never runs. Python evaluates and from left to right and stops at the first false part, so pred(lo) is not called when the list is empty.

Using template 3 for a count

If you know the last yes, you know the count of yes answers, because all the yes answers are at the start. The count of values that are 4 or less is last_true(...) + 1. With the list above, that is 3 + 1, which is 4. When the function returns -1, the count is -1 + 1, which is 0. So the formula works in every case.

A shortcut that uses template 2

You can also find the last true with template 2. Flip the question. The last true is the index just before the first false. So last_true equals first_true of the opposite predicate, minus 1. This is useful when you already have template 2 and do not want a second loop. The main template is still the one above, because it returns -1 directly and reads more clearly in code review.

Putting it into one program

The program in the Template 1 section contains all three templates. The predicates are written inside the calls, so you can change them and run again.

Check two lines by hand. first_true a[i] > 4 is 4, because index 4 holds 7, the first value above 4. last_true a[i] <= 1 is -1, because no value is 1 or less.

Avoiding the infinite loop

A binary search loop is safe when every branch inside it makes the range smaller. If one branch leaves the range the same size, the loop can repeat forever. The rule is simple to check. For each branch, ask: "after this line, is the new range smaller than the old one?"

Check each template.

• Template 1 moves lo to mid + 1 or hi to mid - 1. Both leave mid out of the range, so the range always shrinks.
• Template 2 moves lo to mid + 1, which shrinks the range. It moves hi to mid. This also shrinks the range, because mid is always less than hi when lo is less than hi. The formula lo + (hi - lo) // 2 never reaches hi when lo is less than hi.
• Template 3 moves lo to mid. This shrinks the range only because mid is greater than lo, and that is guaranteed by rounding up. It moves hi to mid - 1, which always shrinks.

Now break template 3 on purpose. Change the middle formula back to the rounding-down version, and move lo to mid as before. Here is the broken loop. It is the same as the last_true template with the rounding changed:

lo = 0
hi = n - 1
while lo < hi:
    mid = lo + (hi - lo) // 2
    if pred(mid):
        lo = mid
    else:
        hi = mid - 1

Take n = 2 and the predicate i <= 0. The first answer is yes and the second is no, so the last true is index 0. Start with lo = 0 and hi = 1. The middle is 0 + (1 - 0) // 2, which is 0. The predicate at 0 is a yes, so lo = mid sets lo to 0 again. Nothing changed. The loop repeats with the same numbers forever.

To show this without hanging the terminal, the copy below has a counter that stops the loop after 20 steps and prints the state. The counter is not part of the real template.

The output is:

  still looping after 20 steps: lo=0 hi=1
result: None

The real loop has no counter, so it would never stop. The fix is the rounding-up formula from the template. With lo = 0 and hi = 1, the middle becomes 0 + (1 - 0 + 1) // 2, which is 1. A yes at index 1 would move lo to 1, and the loop would end.

Common mistakes

Mistake 1: the upper end starts one too far

Template 1 starts hi at len(a) - 1. A common mistake is to write len(a). The loop then reads past the end of the list. Here is a small function with that mistake. It searches for 9 in [1, 3, 5].

def find_exact_off_by_one(a, target):
    lo = 0
    hi = len(a)
    while lo <= hi:
        mid = lo + (hi - lo) // 2
        if a[mid] == target:
            return mid
        if a[mid] < target:
            lo = mid + 1
        else:
            hi = mid - 1
    return -1
 
 
print(find_exact_off_by_one([1, 3, 5], 9))

Python stops the program with this error. The folder part of each file path is shortened to ....

Traceback (most recent call last):
  File ".../mistake1.py", line 15, in <module>
    print(find_exact_off_by_one([1, 3, 5], 9))
          ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^
  File ".../mistake1.py", line 6, in find_exact_off_by_one
    if a[mid] == target:
       ~^^^^^
IndexError: list index out of range

The last line names the error type and says what went wrong: the program asked for an index that the list does not have. The lines above it are the traceback, which lists the function calls that led to the error, with the most recent call last. The ^ marks under a[mid] show the part of the line that failed. The trace shows that mid became 3, because hi started at 3. The list has 3 values, so its valid indexes are 0, 1 and 2. The fix is to start hi at len(a) - 1.

Mistake 2: forgetting the return after the loop

Python does not check for a missing return when you save or run the file. A function that reaches its end without a return gives back None, Python's word for "no value". The mistake shows up later, as a different error, at the place where the result is used. Here is the template without its last line.

def find_exact_no_return(a, target):
    lo = 0
    hi = len(a) - 1
    while lo <= hi:
        mid = lo + (hi - lo) // 2
        if a[mid] == target:
            return mid
        if a[mid] < target:
            lo = mid + 1
        else:
            hi = mid - 1
 
 
result = find_exact_no_return([2, 4, 7], 5)
print("result:", result)
print("result + 1:", result + 1)

The output is:

result: None
Traceback (most recent call last):
  File ".../mistake2.py", line 16, in <module>
    print("result + 1:", result + 1)
                         ~~~~~~~^~~
TypeError: unsupported operand type(s) for +: 'NoneType' and 'int'

The search for 5 fails, so the function should return -1. Instead it returns None, and the addition fails. The fix is to add return -1 after the loop. The error message says NoneType, which is the type of None.

Mistake 3: / gives a decimal number

Python's / always gives a decimal number, even when the answer is a whole number. A list index must be a whole number, so the / version fails. Here is the middle formula with the wrong operator.

a = [2, 4, 4, 4, 7, 9, 12]
lo = 0
hi = 6
mid = lo + (hi - lo) / 2
print("mid:", mid)
print("a[mid]:", a[mid])

The output is:

mid: 3.0
Traceback (most recent call last):
  File ".../mistake3.py", line 6, in <module>
    print("a[mid]:", a[mid])
                     ~^^^^^
TypeError: list indices must be integers or slices, not float

The value 3.0 looks right, but it is a decimal number, so the list refuses it. The fix is //. In Java, / between two whole numbers already drops the fraction, so this mistake does not happen there.

Mistake 4: moving hi to mid minus one in template 2

In template 2, a yes at mid must set hi to mid, not mid - 1. The mistake is easy to make, because Template 1 uses mid - 1. In template 2, the mid index might be the first yes, so mid - 1 throws it away.

def first_true_wrong(n, pred):
    lo = 0
    hi = n
    while lo < hi:
        mid = lo + (hi - lo) // 2
        if pred(mid):
            hi = mid - 1
        else:
            lo = mid + 1
    return lo
 
 
a = [2, 4, 4, 4, 7, 9, 12]
print("first_true_wrong on a >= 4 ->", first_true_wrong(len(a), lambda i: a[i] >= 4), "(correct is 1)")

The output is:

first_true_wrong on a >= 4 -> 0 (correct is 1)

The trace shows why. The middle is 3, a yes, so hi becomes 2. The middle is 1, a yes, so hi becomes 0. The range is now 0 to 0, and the function returns 0. But index 0 holds 2, which is not a yes. The index 1 was removed from the range when hi became 0.

Mistake 5: the middle index overflows in other languages

In Java and C, an int holds only whole numbers up to about 2 billion. A sum above that limit wraps around to a negative number, so (lo + hi) / 2 can go wrong for large lists. Python does not have this limit. An integer grows as large as it needs to, so the sum is always correct.

The output is:

(lo + hi) // 2 = 1750000000
lo + (hi - lo) // 2 = 1750000000
2**70 + 2**70 = 2361183241434822606848

Here both middle formulas agree, and the big sum is printed in full. The 2**70 means 2 to the power 70.

Mistake 6: using -1 as an index

Template 1 returns -1 when the target is missing. Python accepts a negative index, so a[-1] does not fail. It gives the last value in the list. If the caller uses the -1 without checking, the program keeps going with a wrong value and no error.

The output is:

index: -1
a[index]: 12

The value 5 is not in the list, but the program prints 12 without an error. Check the result for -1 before you use it as an index.

Mistake 7: a predicate that does not change from false to true

Each template needs its predicate to switch once. For first_true, the answers must go no and then yes, and never switch back. If the answers switch back and forth, the template returns an index with no guarantee. For example, the predicate lambda x: x == 4 is true only at one point. Searching it with template 2 is wrong, because the answers are no, no, yes, no, no. The template does not know that a yes can be followed by a no.

Before you use a template, write the answers for a few indexes on paper. If the pattern is not one clean switch, a binary search will not work on that question.

Practice

Each exercise has the answer under it. Try the work on paper first.

Exercise 1. Trace first_true with n = 5 and the predicate i >= 3. Write the value of lo, hi and mid after each loop step, and the final result.

Show the answer

Start with lo 0 and hi 5. The middle is 2, and 2 >= 3 is no, so lo becomes 3. The range is now 3 to 5, and the middle is 4. 4 >= 3 is yes, so hi becomes 4. The range is 3 to 4, and the middle is 3. 3 >= 3 is yes, so hi becomes 3. The range is 3 to 3, so the loop ends. The result is 3.

Exercise 2. The list is [1, 2, 2, 2, 3, 5]. Using first_true for the lower bound and the upper bound, find how many copies of 2 are in the list. Then find how many copies of 4 are in the list.

Show the answer

For 2, the lower bound is the first index whose value is at least 2, which is 1. The upper bound is the first index whose value is greater than 2, which is 4. The count is 4 minus 1, which is 3. For 4, the lower bound is the first index whose value is at least 4, which is 5. The upper bound is the first index whose value is greater than 4, which is also 5. The count is 5 minus 5, which is 0.

Exercise 3. Use last_true to find the last index of the list [-3, -1, 0, 2] whose value is below zero. Then say what the function returns when every value is zero or more.

Show the answer

The values below zero are at indexes 0 and 1, so the last one is index 1. The function returns 1. If every value is zero or more, no index is a yes, so the function returns -1.

Exercise 4. The code below is meant to find the target in a sorted list, but it can loop forever. Find the line that causes the problem and explain why.

lo = 0
hi = len(a) - 1
while lo <= hi:
    mid = lo + (hi - lo) // 2
    if a[mid] < target:
        lo = mid
    elif a[mid] > target:
        hi = mid - 1
    else:
        return mid
return -1

Show the answer

The line lo = mid is wrong. It should be lo = mid + 1. Take the list [1, 5] and the target 9. With lo at 0 and hi at 1, the middle is 0. The value at index 0 is 1, which is less than 9, so lo = mid sets lo to 0 again. The range never changes, and the loop repeats forever. The fix is lo = mid + 1, which removes mid from the range.

Exercise 5. Write a function that returns the number of values in a sorted list that are 5 or less. Use last_true from this lesson. Test it on [1, 3, 5, 6, 8], which should return 3, and on [6, 8], which should return 0.

Show the answer

Return last_true(len(a), lambda i: a[i] <= 5) + 1. For [1, 3, 5, 6, 8], the last value that is 5 or less is at index 2, so the result is 2 + 1, which is 3. For [6, 8], last_true returns -1, so the result is -1 + 1, which is 0. The + 1 works because the yes answers always form a block at the start of the list, and the count of that block is the last yes index plus 1.

Summary

• Keep lo and hi as the range that must contain the answer, and write the loop so that every branch shrinks it.
• Use while lo <= hi to find an exact value, and return inside the loop.
• Use while lo < hi with hi = mid for the first true, and return lo after the loop.
• Use while lo < hi with lo = mid and rounding up for the last true, and check pred(lo) after the loop.
• Compute the middle as lo + (hi - lo) // 2, and use // so the index stays a whole number.
• Before writing the code, write the yes and no answers for a few indexes. The answers must switch once, in the right direction, for the template to work.