Binary search on the answer in Python

Why this matters

Many problems ask for "the smallest value that works" or "the largest value that works". Examples are the smallest truck that can carry all parcels within a set number of days, the largest piece length you can cut to get enough pieces, and the smallest speed that finishes a job on time. A first attempt is usually to try every possible value, one by one, and keep the first one that works. That can be far too slow when the range of values is large.

This lesson shows a different method. You do not search through a list. You search through the possible answers. Each guess is checked with a yes-or-no test, and the test tells you which half of the answers to drop. The loop shapes come from the previous lesson in this module, the iterative binary search templates, so this lesson uses them without explaining them again.

The idea: guess an answer and check it

Start with a small example that you can check by hand. Find the smallest whole number x such that x * x is at least n. This is the ceiling of the square root of n. For n = 17, the answer is 5, because 4 times 4 is 16, which is too small, and 5 times 5 is 25, which is big enough.

A simple method tries 0, then 1, then 2, and so on. That works, but for n near one billion it tries about 31,623 numbers. Binary search tries far fewer.

Here is the idea. You know the answer is somewhere between 0 and n. Pick the middle value of that range and ask a yes-or-no question: "does this value work?" For x * x >= n, the answer to "does x work" is yes or no. If the middle value works, the smallest answer is at or below it, so throw away everything above. If the middle value does not work, the answer is above it, so throw away everything at or below it. Each question removes half of the remaining range.

The yes-or-no question is called a check, sometimes written as ok(x). Two conditions must hold before binary search on the answer is correct.

1. There is a range. You know a smallest possible answer and a largest possible answer. The true answer is inside that range.
2. The check is monotone. Monotone means the answers never switch back. For a "smallest value that works" question, once a value works, every larger value also works. So the answers go no, no, no, yes, yes, yes. For a "largest value that works" question, the answers go yes, yes, yes, no, no, no.

If the check has these two properties, you can use the templates from the previous lesson with the check as the argument.

Python has one whole-number type, int. It has no size limit, so x * x can be as large as memory allows. There is no separate long type to choose. A check is a function, and a function can be passed as an argument just like a number. A lambda is a short way to write a small function on one line. For example, lambda x: x * x >= n is a function that takes x and returns True or False.

Here are two functions that wrap the templates. The first finds the smallest value that passes the check. The second finds the largest. The names are least_feasible and largest_feasible.

from collections.abc import Callable
 
 
def least_feasible(lo: int, hi: int, ok: Callable[[int], bool]) -> int:
    while lo < hi:
        mid = lo + (hi - lo) // 2
        if ok(mid):
            hi = mid
        else:
            lo = mid + 1
    return lo
 
 
def largest_feasible(lo: int, hi: int, ok: Callable[[int], bool]) -> int:
    while lo < hi:
        mid = lo + (hi - lo + 1) // 2
        if ok(mid):
            lo = mid
        else:
            hi = mid - 1
    return lo if ok(lo) else lo - 1

least_feasible is the first-true template with the check passed in as ok. largest_feasible is the last-true template. The // operator is floor division. It divides and drops the fraction, so 7 // 2 is 3. The / operator always gives a decimal number, which is a mistake in this code. The Common mistakes section shows what happens.

Two parts of the code are type hints. A type hint is a note after a name that says what kind of value it should hold, such as int or bool. Python does not enforce hints while the program runs. Tools such as mypy read them and report mistakes. Callable[[int], bool] means a function that takes one int and returns a bool.

The end of largest_feasible returns lo - 1 when no value in the range passes the check. Since the smallest value in the range is lo, lo - 1 is one below the range, which is a clear signal that no answer exists.

Example 1: the ceiling of a square root

This is the easiest case, so it shows the whole method with no extra rules. The answer lies between 0 and n. The check is x * x >= n, which is a monotone check, because once x is large enough, every larger x also passes.

def ceil_sqrt(n: int) -> int:
    return least_feasible(0, n, lambda x: x * x >= n)

Take n = 17 and trace the search.

1. The range is 0 to 17. mid is 8. The check says 8 times 8 is 64, which is at least 17, so yes. Set hi to 8.
2. The range is 0 to 8. mid is 4. The check says 16 is at least 17, which is no. Set lo to 5.
3. The range is 5 to 8. mid is 6. The check says 36 is at least 17, which is yes. Set hi to 6.
4. The range is 5 to 6. mid is 5. The check says 25 is at least 17, which is yes. Set hi to 5.
5. The range is 5 to 5. Return 5.

The range is never empty, because x * x >= n is yes at x = n for every n from 1 up. The one special case is n = 0. Then the range is 0 to 0, and the function returns 0 without a check, which is correct.

Python's int has no size limit, so this works for very large numbers too. ceil_sqrt(10 ** 40) returns 100000000000000000000, which is exact. The ** operator means power.

The standard library version

The bisect module is part of the standard library. Its name means "split in two". Since Python 3.10, bisect_left accepts a key function. It searches a sorted sequence for the first place where the key is at least the target. A range is a sorted sequence of whole numbers that stores only its start, stop and step, so it does not build a list in memory. Use True as the target, because False is less than True in Python.

import bisect
 
 
def ceil_sqrt_bisect(n: int) -> int:
    return bisect.bisect_left(range(0, n + 1), True, key=lambda x: x * x >= n)

Calling it with 15, 16 and 17 gives 4, 4 and 5. It works here because the check is monotone and range is a sequence. The hand-written functions above work for any range of whole numbers, and the rest of this lesson uses them.

Example 2: packing parcels into a fixed number of days

Now a problem with a real rule to check. You have parcels in a line. Each parcel has a weight. A truck carries parcels in the order they are given, and it cannot take them out of order. Each day, the truck carries a load up to some capacity. You need to ship every parcel within days days. Find the smallest capacity that makes this possible.

For example, take the weights [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] and 5 days. Try a capacity of 15.

• Day 1 takes 1, 2, 3, 4, and 5. The total is 15, which fits.
• Day 2 takes 6 and 7. The total is 13. Adding 8 would make 21, which is too much.
• Day 3 takes 8. Adding 9 would make 17, which is too much.
• Day 4 takes 9. Adding 10 would make 19, which is too much.
• Day 5 takes 10.

Five days is enough. Now try a capacity of 14. Day 1 takes 1 to 4, which is 10. Adding 5 would make 15, which is too much. The same kind of filling needs six days. So 15 is the smallest capacity that works. The capacity is the answer we want, which makes this a "smallest value that works" problem.

The check

The check for one capacity is a greedy fill. Greedy means that at each step you make the choice that looks best right now. Here the choice is to keep adding parcels to the current day until the next one would go over the capacity. Then start a new day. Count the days used. The capacity passes if the count is at most days.

Greedy is correct here. Filling a day as much as possible never leaves you with more days later than a plan that leaves room unused. So if the greedy count is more than days, no other plan with that capacity can do better.

def can_ship(weights: list[int], days: int, capacity: int) -> bool:
    days_used = 1
    load = 0
    for w in weights:
        if w > capacity:
            return False
        if load + w > capacity:
            days_used += 1
            load = 0
        load += w
    return days_used <= days

list[int] is a type hint for a list of whole numbers. The line if w > capacity: return False is needed. A parcel heavier than the capacity can never be carried, so no number of days would help. Without that line, the loop would put that parcel on a new day with a load bigger than the capacity, and the check would pass wrongly.

The bounds

The smallest possible capacity is the weight of the heaviest parcel. Any smaller capacity cannot carry that parcel. The largest useful capacity is the sum of all weights, because then every parcel fits on one day. The answer is between those two numbers.

def least_capacity(weights: list[int], days: int) -> int:
    lo = max(weights, default=0)
    hi = sum(weights)
    return least_feasible(lo, hi, lambda c: can_ship(weights, days, c))

max and sum are built-in functions. default=0 makes max return 0 for an empty list, instead of raising an error. The lambda uses weights and days from the surrounding function, so the check only needs the capacity as its argument.

Here is the complete program, with the other two examples in it. Press Run. The program needs no input.

The output is:

ceil_sqrt(15) = 4
ceil_sqrt(16) = 4
ceil_sqrt(17) = 5
parcels, 5 days: capacity 15
parcels, 2 days: capacity 28
ribbons, 3 pieces: length 5
ribbons, 4 pieces: length 4
ribbons, 22 pieces: length 0

The parcels line for 5 days matches the hand check above. The line for 2 days is 28, because the best split is 1 to 7 on one day, which is 28, and 8 to 10 on the other day, which is 27.

Why the search is fast

Each check is a loop over every parcel, so it costs time proportional to the number of parcels. The search makes about as many checks as it takes to halve the range down to one value. Here is a run with a bigger input: 100,000 parcels whose weights are 1, 2, 3, and so on up to 100,000, and 30 days. The range held 4,999,950,001 values, and each check halves it, so the run made 32 checks. The answer for that run was a capacity of 166,701,564. A plain search that tried each capacity would have needed up to about 5 billion checks, each one a loop over 100,000 parcels.

Example 3: cutting pieces of the same length

The third example reverses the direction. You have ribbons of different lengths, measured in whole units. You need at least k pieces, all the same whole length. Each ribbon can be cut into as many whole pieces as fit, and any leftover is wasted. Find the largest piece length that gives at least k pieces.

Take the ribbons [9, 7, 5] and k = 3. Try a length of 5. The 9 gives one piece with 4 left over, the 7 gives one piece with 2 left over, and the 5 gives one piece. That is 3 pieces, which is enough. Try a length of 6. The 9 gives one piece, the 7 gives one piece, and the 5 gives none. That is 2 pieces, which is too few. So 5 is the answer.

The check is a count. For a piece length L, add up length // L over all ribbons. The // operator rounds down, so 9 // 5 is 1. Pass if the total is at least k.

length // piece for length in lengths is a generator expression. It produces one value at a time, and sum adds them up. It does the same work as a loop with a running total, in one line.

def can_cut(lengths: list[int], k: int, piece: int) -> bool:
    count = sum(length // piece for length in lengths)
    return count >= k
 
 
def longest_piece(lengths: list[int], k: int) -> int:
    return largest_feasible(1, max(lengths), lambda piece: can_cut(lengths, k, piece))

As the piece length gets bigger, the count gets smaller, or stays the same. So the check goes yes, yes, yes, no, no, no. That is the last-true pattern, so largest_feasible is the right template.

The range starts at 1, not 0. A piece length of 0 would divide by zero, so 1 is the smallest length that makes sense. The top of the range is the longest ribbon, because a piece cannot be longer than the ribbon it comes from.

Trace k = 3 with the range 1 to 9.

1. mid is 1 + (9 - 1 + 1) // 2, which is 5. The count at 5 is 1 + 1 + 1, which is 3. That is yes. Set lo to 5.
2. The range is 5 to 9. mid is 5 + (9 - 5 + 1) // 2, which is 7. The count at 7 is 1 + 1 + 0, which is 2. That is less than 3, so no. Set hi to 6.
3. The range is 5 to 6. mid is 5 + (6 - 5 + 1) // 2, which is 6. The count at 6 is 1 + 1 + 0, which is 2. That is no. Set hi to 5.
4. The range is 5 to 5, so the loop ends. The check at 5 is yes, so return 5.

If the total length of all ribbons is less than k, no length works. For example, with k = 22 and a total length of 21, every check fails. largest_feasible then returns lo - 1, which is 0. The caller must treat 0 as "no answer", because a real piece length is at least 1.

Common mistakes

Mistake 1: the top of the range is not a valid answer

least_feasible assumes that the top of the range passes the check. If it does not, the function returns the top with no error. Here the top is 10, the heaviest parcel, instead of the sum, 55. With 2 days, no capacity up to 10 can work, because the truck would need a capacity of at least 28.

parcels = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
answer = least_feasible(1, 10, lambda c: can_ship(parcels, 2, c))
print(f"bad upper bound, 2 days, answer = {answer}, can_ship at answer = {can_ship(parcels, 2, answer)}")

Output:

bad upper bound, 2 days, answer = 10, can_ship at answer = False

The output shows the bug. The answer 10 is returned, and the check at 10 says no. The fix is to set the top of the range to a value that is known to pass, such as the sum of all weights. A good habit is to test the check at the answer you got. If the check fails there, the bounds or the check are wrong.

Mistake 2: the check goes both ways

The method only works when the check changes once. If the check is yes, then no, then yes, the answer can be wrong.

print(least_feasible(0, 10, lambda x: x == 4))

Output:

10

The true answer is 4, but the function returns 10. The check never says yes at the midpoints the search visits, so the search moves to the top. Before you use this function, write the check's answers for a few values and confirm they switch once.

Mistake 3: using / instead of //

Java's int type wraps around when a sum goes past its largest value. Python's int does not wrap around, so that mistake cannot happen here. The nearest mistake in Python is the division operator. / always gives a float, which is a number with a decimal point. A float is not exact for every whole number, and it cannot be used as a list index.

Here is the search with / in the middle line. The only change is / 2 in place of // 2.

def least_feasible_slash(lo: int, hi: int, ok: Callable[[int], bool]) -> int:
    while lo < hi:
        mid = lo + (hi - lo) / 2
        if ok(mid):
            hi = mid
        else:
            lo = mid + 1
    return lo
 
 
print(least_feasible_slash(0, 17, lambda x: x * x >= 17))

Output:

4.6875

The answer should be 5. The function returns a float with no error message, so the bug is easy to miss. If the float is used as a list index, Python stops with an error. The last line of that error is:

TypeError: list indices must be integers or slices, not float

Use // for every midpoint in a binary search.

Mistake 4: a strict comparison changes the answer

For the ceiling of a square root, the check is x * x >= n. If you write x * x > n by mistake, the answer changes. For n = 16, the correct answer is 4, because 4 times 4 is exactly 16. The check with > says 4 does not pass, so the search returns 5.

print(least_feasible(0, 16, lambda x: x * x > 16))

Output:

5

Check the boundary value by hand. When the question says "at least", use >=.

Practice

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

Exercise 1. Write a function that returns the square root rounded down, which is the largest x such that x * x is at most n. Use largest_feasible. What does it return for n = 17 and for n = 16?

Show the answer

Use largest_feasible(0, n, lambda x: x * x <= n). For n = 17, the result is 4, because 4 times 4 is 16, which is at most 17, and 5 times 5 is 25, which is not. For n = 16, the result is 4, because 4 times 4 is exactly 16. The standard library also has math.isqrt(n), which returns the same value. math.isqrt(17) is 4, and math.isqrt(16) is 4.

Exercise 2. Five parcels each weigh 4, in this order: [4, 4, 4, 4, 4]. The truck must ship them within 3 days. What is the smallest capacity that works?

Show the answer

The answer is 8. With a capacity of 8, the first day takes two parcels (8), the second day takes two (8), and the third day takes the last one. That is 3 days. With a capacity of 7, each day can take only one parcel, because two would make 8, which is too much. Five parcels then need 5 days, which is more than 3.

Exercise 3. Ribbons have lengths [10, 10]. You need 5 pieces of the same whole length. What is the longest length that works?

Show the answer

The answer is 3. With a length of 4, each 10 gives two pieces, for a total of 4, which is less than 5. With a length of 3, each 10 gives three pieces, for a total of 6, which is at least 5.

Exercise 4. Remove the line if w > capacity: return False from can_ship. Then call it with the parcels [5, 1], a capacity of 3, and 5 days. What does it return, and what should it return?

Show the answer

Without the line, it returns True. The parcel of 5 is put on a day by itself with a load of 5, which is over the capacity of 3, and the check does not notice. It should return False, because a parcel of 5 can never fit in a truck with a capacity of 3, however many days are available. With the line in place, the check returns False as it should.

Exercise 5. Write ceil_sqrt using a check that is wrong: lambda x: x * x > n. What does it return for n = 16, and why does this differ from the correct answer?

Show the answer

It returns 5. The check asks for a square strictly greater than 16. The square of 4 is exactly 16, so 4 fails the check, and the smallest value that passes is 5. The correct check is x * x >= n, which accepts 4 because 16 is at least 16.

Summary

• A binary search on the answer searches over possible answers, not over a list.
• The check must be monotone, so its answers switch only once across the range.
• Use least_feasible for the smallest value that works, and largest_feasible for the largest value that works.
• Put a bound on each end of the range, and make sure the top of the range is a value that passes.
• Use // for midpoints and whole-number division. A / gives a float, which can give a wrong answer without an error.