{"id":"dsa-complexity-binary-search","slug":"binary-search","title":"Binary search","text":"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:\n\n```python\ndef find(values, target):\n    lo = 0\n    hi = len(values) - 1\n    while lo <= hi:\n        mid = (lo + hi) // 2\n        print(\"looks at index\", mid, \"value\", values[mid])\n        if values[mid] == target:\n            return mid\n        if values[mid] < target:\n            lo = mid + 1\n        else:\n            hi = mid - 1\n    return -1\n\n\nvalues = [3, 8, 14, 21, 30, 42, 57, 61, 77]\nprint(find(values, 57))\nprint(find(values, 10))\n```\n\n```output\nlooks at index 4 value 30\nlooks at index 6 value 57\n6\nlooks at index 4 value 30\nlooks at index 1 value 8\nlooks at index 2 value 14\n-1\n```\n\n**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.\n\n**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.\n\n### The off-by-one traps\n\nEvery `+ 1` and `- 1` matters:\n\n- With `while lo < hi`, the loop stops when one item is left, without looking at it. `find(values, 77)` would return -1.\n- 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.\n- With `hi = len(values)`, the range reaches one place past the end of the list, and a search for a big value looks there:\n\n```python\ndef find(values, target):\n    lo = 0\n    hi = len(values)\n    while lo <= hi:\n        mid = (lo + hi) // 2\n        if values[mid] == target:\n            return mid\n        if values[mid] < target:\n            lo = mid + 1\n        else:\n            hi = mid - 1\n    return -1\n\n\nprint(find([3, 8, 14], 99))\n```\n\n```crash\nTraceback (most recent call last):\n  File \"main.py\", line 15, in <module>\n    print(find([3, 8, 14], 99))\n          ~~~~^^^^^^^^^^^^^^^^\n  File \"main.py\", line 6, in find\n    if values[mid] == target:\n       ~~~~~~^^^^^\nIndexError: list index out of range\n```\n\n### Where would it go?\n\nOften 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`:\n\n```python\ndef first_at_least(values, target):\n    lo = 0\n    hi = len(values)\n    while lo < hi:\n        mid = (lo + hi) // 2\n        if values[mid] < target:\n            lo = mid + 1\n        else:\n            hi = mid\n    return lo\n```\n\nBoth 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`.\n\n### The `bisect` module\n\nPython 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`:\n\n```pycon\n>>> from bisect import bisect_left, bisect_right\n>>> prices = [5, 8, 8, 12, 20, 20, 20, 35, 50]\n>>> bisect_left(prices, 20)\n4\n>>> bisect_right(prices, 20)\n7\n>>> bisect_right(prices, 20) - bisect_left(prices, 20)\n3\n>>> bisect_left(prices, 10)\n3\n>>> bisect_right(prices, 35) - bisect_left(prices, 8)\n7\n```\n\nA 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.\n\n### Searching for the answer\n\nBinary 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]`:\n\n```python\ndef days_needed(weights, capacity):\n    days = 1\n    load = 0\n    for weight in weights:\n        if load + weight > capacity:\n            days += 1\n            load = 0\n        load += weight\n    return days\n\n\ndef least_capacity(weights, days):\n    lo = max(weights)\n    hi = sum(weights)\n    while lo < hi:\n        mid = (lo + hi) // 2\n        if days_needed(weights, mid) <= days:\n            hi = mid\n        else:\n            lo = mid + 1\n    return lo\n\n\nweights = [3, 2, 2, 4, 1, 4]\nprint([days_needed(weights, capacity) for capacity in range(4, 9)])\nprint(least_capacity(weights, 3))\n```\n\n```output\n[5, 4, 3, 3, 3]\n6\n```\n\nThe 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.\n","task":"What does this program print? Follow `lo` and `hi` round by round. Type each line it prints.\n","kind":"predict","seed":"def find(values, target):\n    lo = 0\n    hi = len(values) - 1\n    while lo <= hi:\n        mid = (lo + hi) // 2\n        print(values[mid])\n        if values[mid] == target:\n            return mid\n        if values[mid] < target:\n            lo = mid + 1\n        else:\n            hi = mid - 1\n    return -1\n\n\nprint(find([2, 5, 9, 12, 20, 25, 31, 40, 48, 50, 63], 45))\n","solution":"def find(values, target):\n    lo = 0\n    hi = len(values) - 1\n    while lo <= hi:\n        mid = (lo + hi) // 2\n        print(values[mid])\n        if values[mid] == target:\n            return mid\n        if values[mid] < target:\n            lo = mid + 1\n        else:\n            hi = mid - 1\n    return -1\n\n\nprint(find([2, 5, 9, 12, 20, 25, 31, 40, 48, 50, 63], 45))\n","hints":["The list has 11 values, at indexes 0 to 10, so the first `mid` is (0 + 10) // 2 = 5, the value 25.","25 is smaller than 45, so `lo` becomes 6. The next `mid` is (6 + 10) // 2 = 8, the value 48, which is greater: `hi` becomes 7.","Keep going until `lo` is greater than `hi`. 45 isn't in the list, so the last line is -1."],"tests":[{"name":"Output","stdin":"","expect":"25\n48\n31\n40\n-1","hidden":false}],"require":[],"forbid":[],"lines":["25","48","31","40","-1"],"more":[{"task":"This binary search should return the index of `target` in the sorted list `values`, or -1 if it isn't there. It misses some values that are in the list. Fix it, and keep it a binary search: the hidden test makes about 10,000 searches in a list of 1,000,000 values, too many for a search that goes through the list one by one.\n\nFor example, `find([4, 9, 15], 15)` should return:\n\n```\n2\n```\n","kind":"code","seed":"def find(values, target):\n    lo = 0\n    hi = len(values) - 1\n    while lo < hi:\n        mid = (lo + hi) // 2\n        if values[mid] == target:\n            return mid\n        if values[mid] < target:\n            lo = mid + 1\n        else:\n            hi = mid - 1\n    return -1\n\n\nif __name__ == \"__main__\":\n    print(find([4, 9, 15], 15))\n","solution":"def find(values, target):\n    lo = 0\n    hi = len(values) - 1\n    while lo <= hi:\n        mid = (lo + hi) // 2\n        if values[mid] == target:\n            return mid\n        if values[mid] < target:\n            lo = mid + 1\n        else:\n            hi = mid - 1\n    return -1\n\n\nif __name__ == \"__main__\":\n    print(find([4, 9, 15], 15))\n","hints":["Trace `find([4, 9, 15], 15)`: `lo` is 0 and `hi` is 2, `mid` is 1, and 9 is smaller than 15, so `lo` becomes 2. Is the loop's condition still true?","`lo` and `hi` are both inside the range, so when they are equal, one item is still left to look at.","The loop must also run when `lo == hi`: `while lo <= hi`."],"tests":[{"name":"find([4, 9, 15], 15)","stdin":"","call":"find([4, 9, 15], 15)\n","expect":"2","hidden":false,"returns":true},{"name":"find([4, 9, 15], 4)","stdin":"","call":"find([4, 9, 15], 4)\n","expect":"0","hidden":true,"returns":true},{"name":"find([7], 7)","stdin":"","call":"find([7], 7)\n","expect":"0","hidden":true,"returns":true},{"name":"find([], 3)","stdin":"","call":"find([], 3)\n","expect":"-1","hidden":true,"returns":true},{"name":"[find([1, 3, 5, 7, 9, 11], x) for x in [1, 5, 6, 11, 12, 0]]","stdin":"","call":"[find([1, 3, 5, 7, 9, 11], x) for x in [1, 5, 6, 11, 12, 0]]\n","expect":"[0, 2, -1, 5, -1, -1]","hidden":true,"returns":true},{"name":"Test 6","stdin":"","call":"values = list(range(0, 3_000_000, 3))\nprint(sum(find(values, x) for x in range(0, 3_000_000, 301)))\n","expect":"1661363759","hidden":true,"big":true}],"require":[],"forbid":[]},{"task":"Write `count_in_ranges(values, queries)`. The list `values` is in no particular order, and each query is a tuple `(low, high)`, where `low` is never greater than `high`. A query's answer is how many values are at least `low` and at most `high`. Return a list of the answers, in the order of the queries. The hidden test has 20,000 queries about a list of 300,000 values: a loop over the list for each query is too slow, and so is sorting the list again for each one.\n\nFor example, `count_in_ranges([21, 15, 30, 12, 15, 18, 15], [(15, 21), (19, 29), (0, 100)])` returns:\n\n```\n[5, 1, 7]\n```\n","kind":"code","seed":"def count_in_ranges(values, queries):\n    pass\n\n\nif __name__ == \"__main__\":\n    print(count_in_ranges([21, 15, 30, 12, 15, 18, 15], [(15, 21), (19, 29), (0, 100)]))\n","solution":"from bisect import bisect_left, bisect_right\n\n\ndef count_in_ranges(values, queries):\n    ordered = sorted(values)\n    answers = []\n    for low, high in queries:\n        answers.append(bisect_right(ordered, high) - bisect_left(ordered, low))\n    return answers\n\n\nif __name__ == \"__main__\":\n    print(count_in_ranges([21, 15, 30, 12, 15, 18, 15], [(15, 21), (19, 29), (0, 100)]))\n","hints":["Binary search needs a sorted list. Sort once, before the queries: `ordered = sorted(values)` costs O(n log n), and then every query can use two O(log n) searches in `ordered`.","In the sorted list, `bisect_left(ordered, low)` is the index of the first value that is at least `low`, and `bisect_right(ordered, high)` is the index just after the last value that is at most `high`.","A query's answer is the difference of those two indexes. A loop over the queries can unpack each tuple: `for low, high in queries:`."],"tests":[{"name":"count_in_ranges([21, 15, 30, 12, 15, 18, 15], [(15, 21), (19, 29), (0, 100)])","stdin":"","call":"count_in_ranges([21, 15, 30, 12, 15, 18, 15], [(15, 21), (19, 29), (0, 100)])\n","expect":"[5, 1, 7]","hidden":false,"returns":true},{"name":"count_in_ranges([3, 1, 2], [(5, 9), (-5, 1)])","stdin":"","call":"count_in_ranges([3, 1, 2], [(5, 9), (-5, 1)])\n","expect":"[0, 1]","hidden":true,"returns":true},{"name":"count_in_ranges([4, 4, 4], [(4, 4), (5, 5)])","stdin":"","call":"count_in_ranges([4, 4, 4], [(4, 4), (5, 5)])\n","expect":"[3, 0]","hidden":true,"returns":true},{"name":"count_in_ranges([], [(1, 5)])","stdin":"","call":"count_in_ranges([], [(1, 5)])\n","expect":"[0]","hidden":true,"returns":true},{"name":"count_in_ranges([10, 2, 8, 6, 4], [(3, 9), (2, 2), (11, 20)])","stdin":"","call":"count_in_ranges([10, 2, 8, 6, 4], [(3, 9), (2, 2), (11, 20)])\n","expect":"[3, 1, 0]","hidden":true,"returns":true},{"name":"count_in_ranges([5, 1], [])","stdin":"","call":"count_in_ranges([5, 1], [])\n","expect":"[]","hidden":true,"returns":true},{"name":"sum(count_in_ranges([i * 7_919 % 200_003 for i in range(300_000)], [(q, q + 50_000) for q in range(0, 200_000, 10)]))","stdin":"","call":"sum(count_in_ranges([i * 7_919 % 200_003 for i in range(300_000)], [(q, q + 50_000) for q in range(0, 200_000, 10)]))\n","expect":"1312535941","hidden":true,"big":true,"returns":true}],"require":[],"forbid":[]}]}