Built-in helpers, itertools, bisect, and functools¶

~25 min · prerequisite: Deques and heaps

What you will be able to do¶

Use zip, enumerate, and max with a key, and know how each handles uneven inputs.
Use itertools for running totals and combinations, and bisect to count values in a sorted list.
Use functools.cache to remember results, and the exact integer functions in math.

The idea¶

Python ships with helpers that say what you mean in one line, and they are often faster than a hand-written loop. Reach for them before you write the loop yourself. Each helper has a short, precise definition in the library reference, so you can check its behaviour there. The built-in zip(a, b) pairs items by position and stops at the shorter input, so extra items are dropped silently. The built-in enumerate(a, start=1) pairs each item with a counter. The built-in max(a, key=f) returns the item with the largest key, and on a tie it returns the first such item.

The module itertools provides iterators, which produce items one at a time, so you usually wrap them in list(...) to see the values. The function accumulate gives running totals, which are prefix sums. The function combinations(a, 2) gives each pair of positions once. The function groupby groups only neighbouring equal items, so sort the data first.

The module bisect searches a sorted list. The call bisect_left(a, x) returns the position before any items equal to x, and bisect_right(a, x) returns the position after them. Both take logarithmic time. The difference of two calls counts the values in a range. bisect does not check the order of the list, so on unsorted input it gives wrong answers without any error.

The module functools has cache, which remembers the results of a function, so a recursive function does not repeat work. The module math has isqrt, which gives the exact integer square root. Prefer it to int(n ** 0.5), because floating-point values can be off for large numbers.

Worked example¶

from itertools import accumulate
from bisect import bisect_left, bisect_right

values = [3, 1, 4, 1, 5]
print(list(accumulate(values)))

ordered = sorted(values)
lo = bisect_left(ordered, 1)
hi = bisect_right(ordered, 4)
print(ordered[lo:hi])

The output is:

[3, 4, 8, 9, 14]
[1, 1, 3, 4]

Line by line, this is what Python does:

accumulate(values) yields 3, then 3 + 1, then 3 + 1 + 4, and so on. list collects the running sums.
sorted(values) makes the sorted list [1, 1, 3, 4, 5] that bisect needs.
bisect_left(ordered, 1) returns 0, the position before the first 1, and bisect_right(ordered, 4) returns 4, the position after the 4.
The slice ordered[lo:hi] takes positions 0 to 3, which hold the values from 1 to 4.

Try it¶

The first block pairs items. The second block counts values in a range and uses cache and the math functions.

from bisect import bisect_left, bisect_right

def pair_up(a, b):
    return list(zip(a, b))

assert pair_up([1, 2, 3], "abc") == [(1, "a"), (2, "b"), (3, "c")]
assert pair_up([1, 2, 3], [9]) == [(1, 9)]
assert list(enumerate("xy", start=1)) == [(1, "x"), (2, "y")]

def count_in_range(sorted_values, lo, hi):
    return bisect_right(sorted_values, hi) - bisect_left(sorted_values, lo)

assert count_in_range([1, 3, 3, 5, 8, 9], 3, 8) == 4
assert count_in_range([1, 3, 3, 5, 8, 9], 10, 20) == 0

from functools import cache
from math import comb, gcd, isqrt, lcm

@cache
def fib(n):
    return n if n < 2 else fib(n - 1) + fib(n - 2)

assert fib(50) == 12586269025
assert comb(5, 2) == 10
assert gcd(12, 18) == 6
assert lcm(4, 6) == 12
assert isqrt(99) == 9
assert isqrt(10 ** 16) == 10 ** 8

Common mistakes¶

Expecting zip to complain about length. The call list(zip([1, 2, 3], [4])) returns [(1, 4)] with no message. Add strict=True to get:
```
ValueError: zip() argument 2 is shorter than argument 1
```
Grouping unsorted data. The call [(k, len(list(g))) for k, g in groupby([1, 2, 1])] returns [(1, 1), (2, 1), (1, 1)], so the key 1 appears twice. Call sorted(...) first.
Searching an unsorted list with bisect. The call bisect_left([3, 1, 2], 1) returns 0, although the 1 is at index 1. No error appears. Sort the list first.
Passing a list to a cached function. Calling a function decorated with @cache as f([1]) gives:
```
TypeError: unhashable type: 'list'
```
Pass a tuple, such as f((1,)).

Practice¶

Write the function named in each problem. The samples are checks.

Running totals (easy). Write running_totals(values), which returns the prefix sums. Sample: running_totals([3, 1, 4]) returns [3, 4, 8].
Index of the largest (easy). Write index_of_max(values). Return the first index of the largest value. Sample: index_of_max([2, 9, 4, 9]) returns 1.
Count in a range (medium). Write count_in_range(values, lo, hi). The input is not sorted. Count the values v with lo <= v <= hi. Sample: count_in_range([8, 1, 5, 3, 9, 3], 3, 8) returns 4.
Even products (medium). Write even_pairs(values). Count the pairs of positions i < j for which values[i] * values[j] is even. Sample: even_pairs([1, 2, 3, 4]) returns 5.
Sum of gcd values (stretch). Write gcd_sum(n). Return the sum of gcd(i, j) over all pairs 1 <= i < j <= n. Samples: gcd_sum(3) returns 3, and gcd_sum(4) returns 7.

Solution

```python from itertools import accumulate

def running_totals(values): return list(accumulate(values))

assert running_totals([3, 1, 4]) == [3, 4, 8] ```

```python def index_of_max(values): return max(range(len(values)), key=lambda i: values[i])

assert index_of_max([2, 9, 4, 9]) == 1 ```

max returns the first item with the largest key, so the earlier index wins a tie.

```python from bisect import bisect_left, bisect_right

def count_in_range(values, lo, hi): ordered = sorted(values) return bisect_right(ordered, hi) - bisect_left(ordered, lo)

assert count_in_range([8, 1, 5, 3, 9, 3], 3, 8) == 4 ```

```python from itertools import combinations

def even_pairs(values): return sum(1 for a, b in combinations(values, 2) if (a * b) % 2 == 0)

assert even_pairs([1, 2, 3, 4]) == 5 ```

```python from itertools import combinations from math import gcd

def gcd_sum(n): return sum(gcd(i, j) for i, j in combinations(range(1, n + 1), 2))

assert gcd_sum(3) == 3 assert gcd_sum(4) == 7 ```

This checks every pair, so the work grows with the square of n. It is a correct first version.

Recap¶

zip stops at the shorter input, enumerate adds positions, and max returns the first item among ties.
itertools returns iterators, so wrap them in list(...) to see values. bisect needs a sorted list.
functools.cache memoises a function with hashable arguments, and math.isqrt gives an exact integer square root.