Binary search

You met binary search in the recursion module: look at the middle item of a sorted list, and drop the half where the target can't be. Interviews usually want it as a loop. It needs no recursion, only two variables, lo and hi, the edges of the part still to search. This one prints what it looks at:

def find(values, target):
    lo = 0
    hi = len(values) - 1
    while lo <= hi:
        mid = (lo + hi) // 2
        print("looks at index", mid, "value", values[mid])
        if values[mid] == target:
            return mid
        if values[mid] < target:
            lo = mid + 1
        else:
            hi = mid - 1
    return -1


values = [3, 8, 14, 21, 30, 42, 57, 61, 77]
print(find(values, 57))
print(find(values, 10))

looks at index 4 value 30
looks at index 6 value 57
6
looks at index 4 value 30
looks at index 1 value 8
looks at index 2 value 14
-1

Why it's right: the loop keeps one promise: if the target is in the list, it's somewhere from lo to hi. When values[mid] is smaller than the target, so is every value before it, since the list is sorted, and the target can only be after mid: lo = mid + 1. When it's greater, hi = mid - 1. Each round also drops mid itself, so the range shrinks and the loop ends. Once lo passes hi, the range is empty, and the promise says the target isn't there.

What it costs: each round halves the range, so it's O(log n) time: 20 looks for a million values, 30 for a billion. The memory is O(1). The list must be sorted, though. Sorting costs O(n log n), which pays when you search the same list many times.

The off-by-one traps

Every + 1 and - 1 matters:

• With while lo < hi, the loop stops when one item is left, without looking at it. find(values, 77) would return -1.
• With lo = mid instead of mid + 1, once lo and hi are neighbors, mid is lo. If the target is greater, lo = mid changes nothing, and the loop never ends.
• With hi = len(values), the range reaches one place past the end of the list, and a search for a big value looks there:

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


print(find([3, 8, 14], 99))

Traceback (most recent call last):
  File "main.py", line 15, in <module>
    print(find([3, 8, 14], 99))
          ~~~~^^^^^^^^^^^^^^^^
  File "main.py", line 6, in find
    if values[mid] == target:
       ~~~~~~^^^^^
IndexError: list index out of range

Where would it go?

Often you want the place where a value would go: the first index whose value is at least the target. That answer can be "after everything", so hi starts at len(values), one past the end, and the range is from lo up to but not including hi. Since mid may be the answer, it stays in the range, hi = mid, and the loop runs while the range isn't empty, lo < hi:

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

Both styles work. Mixing them causes the first and the last bug above. The middle one is a bug in both styles: hi may stay at mid, but lo must always move past it, lo = mid + 1.

The bisect module

Python has this search ready-made in the bisect module, O(log n) and written in C. bisect_left(values, x) is first_at_least. bisect_right(values, x) is the first index whose value is greater than x. Between the two sit exactly the copies of x:

>>> from bisect import bisect_left, bisect_right
>>> prices = [5, 8, 8, 12, 20, 20, 20, 35, 50]
>>> bisect_left(prices, 20)
4
>>> bisect_right(prices, 20)
7
>>> bisect_right(prices, 20) - bisect_left(prices, 20)
3
>>> bisect_left(prices, 10)
3
>>> bisect_right(prices, 35) - bisect_left(prices, 8)
7

A shop's prices, sorted: 3 items cost 20. bisect_left(prices, 10) is 3, which says two things at once: a 10 would go at index 3, and 3 items cost less than 10, since everything before that index is smaller. The last line counts the items that cost from 8 to 35: 7 of them. In general, the number of values from low to high is bisect_right(values, high) - bisect_left(values, low): two O(log n) searches, however many values there are between. bisect.insort(values, x) puts x in its place, finding it in O(log n), but inserting it is O(n), since every item after it moves.

Searching for the answer

Binary search works on more than lists. Some questions ask for the smallest number that passes a check. A truck ships parcels in the order they're listed, as many each day as fit: what's the smallest capacity that ships them all in days days? Checking one capacity is easy: load the parcels in order, start a new day when the next one doesn't fit, and count the days. And the check flips only once: a capacity that's too small fails, and every capacity from the answer up passes. So binary-search the capacity with the loop of first_at_least, calling the check at mid where it read values[mid]:

def days_needed(weights, capacity):
    days = 1
    load = 0
    for weight in weights:
        if load + weight > capacity:
            days += 1
            load = 0
        load += weight
    return days


def least_capacity(weights, days):
    lo = max(weights)
    hi = sum(weights)
    while lo < hi:
        mid = (lo + hi) // 2
        if days_needed(weights, mid) <= days:
            hi = mid
        else:
            lo = mid + 1
    return lo


weights = [3, 2, 2, 4, 1, 4]
print([days_needed(weights, capacity) for capacity in range(4, 9)])
print(least_capacity(weights, 3))

[5, 4, 3, 3, 3]
6

The capacities 4 to 8 need 5, 4, 3, 3 and 3 days: for 3 days, the check fails for 4 and 5 and passes from 6 on. The capacity is at least the heaviest parcel and at most all of them together, so those are lo and hi. Each check is O(n), and the search makes O(log S) of them, where S is the total weight: O(n log S) in all, where trying every capacity would take O(n·S). This is binary search on the answer. Look for it when a question asks for the smallest or the largest value for which something is possible.