← Back to the 16 Patterns

The Python Toolkit

Five modules, one complexity table, and a dozen traps. This is the part of Python that decides whether your correct idea also runs fast enough.

Interviewers do not test Python trivia. But they do watch you reach for list.pop(0) and quietly note that your O(n) algorithm just became O(n²). This page is the small set of facts that stops that happening.

What each operation actually costs

Know this table. Most accidental time-limit failures are one row of it.

collections

Three classes carry almost all the weight.

Counter

from collections import Counter

counts = Counter("mississippi")
counts["s"]                     # 4
counts["z"]                     # 0, and no KeyError
counts.most_common(2)           # [('i', 4), ('s', 4)]

# Counters compare and combine directly, which anagram problems love.
Counter("listen") == Counter("silent")      # True
Counter("aab") - Counter("ab")              # Counter({'a': 1})

Two behaviours matter. A missing key reads as 0 rather than raising, and most_common(k) uses heapq.nlargest internally, so it is O(n log k) rather than a full sort.

defaultdict

from collections import defaultdict

graph: defaultdict[str, list[str]] = defaultdict(list)
graph["a"].append("b")          # no "if key not in graph" needed

tally: defaultdict[str, int] = defaultdict(int)
tally["x"] += 1

buckets: defaultdict[str, set[str]] = defaultdict(set)
buckets["k"].add("v")

deque

from collections import deque

queue = deque([1, 2, 3])
queue.popleft()                 # O(1). list.pop(0) would be O(n).
queue.appendleft(0)
queue.extendleft([9, 8])        # note: inserts in reverse order

recent = deque(maxlen=3)        # fixed size; pushing past 3 drops from the far end

Use a deque for every BFS queue, for sliding-window maximum, and any time you need both ends. It is a doubly linked list of blocks, so indexing the middle is O(n). If you index it a lot, you wanted a list.

heapq

import heapq

heap = [5, 1, 4]
heapq.heapify(heap)             # O(n), in place
heapq.heappush(heap, 0)
smallest = heapq.heappop(heap)  # 0

heapq.heapreplace(heap, 9)      # pop then push, one sift instead of two
heapq.nlargest(2, [5, 1, 4])    # [5, 4]
heapq.nsmallest(2, [5, 1, 4])   # [1, 4]

bisect

import bisect

data = [1, 3, 3, 5]
bisect.bisect_left(data, 3)     # 1, first index where data[i] >= 3
bisect.bisect_right(data, 3)    # 3, first index where data[i] >  3
bisect.bisect_left(data, 4)     # 3, the insertion point for a missing value

bisect.insort(data, 4)          # data becomes [1, 3, 3, 4, 5]

Since Python 3.10 both accept a key= argument, which removes the old trick of building a parallel list of keys.

itertools

from itertools import accumulate, combinations, permutations, product, groupby, pairwise

list(accumulate([1, 2, 3, 4]))            # [1, 3, 6, 10]  prefix sums
list(accumulate([3, 1, 4], max))          # [3, 3, 4]      running maximum

list(combinations([1, 2, 3], 2))          # [(1, 2), (1, 3), (2, 3)]
list(permutations([1, 2]))                # [(1, 2), (2, 1)]
list(product([0, 1], repeat=2))           # [(0,0), (0,1), (1,0), (1,1)]

list(pairwise([1, 2, 3]))                 # [(1, 2), (2, 3)]  adjacent pairs
[(key, len(list(run))) for key, run in groupby("aaabbc")]   # [('a',3), ('b',2), ('c',1)]

functools

from functools import cache, lru_cache, reduce, cmp_to_key


@cache
def fib(n: int) -> int:
    """Without the decorator this is 2**n calls. With it, n."""
    return n if n < 2 else fib(n - 1) + fib(n - 2)


fib(100)                                   # instant
reduce(lambda a, b: a * b, [1, 2, 3, 4])   # 24

lru_cache(maxsize=None) is the same thing on older versions. cmp_to_key converts an old-style comparison function into a sort key, which is occasionally the cleanest way to express a custom ordering such as Largest Number.

The traps

Mutable default argument

def add(item, bucket=[]):
    # The list is created ONCE, when
    # the function is defined. Every
    # call shares it.
    bucket.append(item)
    return bucket

def add(item, bucket=None):
    if bucket is None:
        bucket = []
    bucket.append(item)
    return bucket

Aliased rows in a 2D grid

grid = [[0] * 3] * 2
grid[0][0] = 1
# grid is now [[1,0,0], [1,0,0]]
# because both rows are the SAME
# list object.

grid = [[0] * 3 for _ in range(2)]
grid[0][0] = 1
# grid is [[1,0,0], [0,0,0]]
# The comprehension builds a fresh
# row each time.

The rest

Idioms worth having

from collections import Counter

# Iterate with an index, or two sequences together.
for index, value in enumerate("abc"):
    pass

for a, b in zip([1, 2], "xy"):
    pass

# Sort by a computed key. sorted() returns a new list; .sort() is in place.
words = ["bbb", "a", "cc"]
by_length = sorted(words, key=len)                 # ['a', 'cc', 'bbb']
by_two = sorted(words, key=lambda w: (len(w), w))  # length, then alphabetical

# Descending without reversing afterwards.
descending = sorted([3, 1, 2], reverse=True)

# max and min take a key too.
longest = max(words, key=len)

# Unpacking, including the starred form.
first, *rest = [1, 2, 3]
head, *middle, tail = [1, 2, 3, 4]

# Swap without a temporary. The right side is evaluated first.
a, b = 1, 2
a, b = b, a

# Chained comparison, which reads exactly like the maths.
n = 5
in_range = 0 <= n < 10

# any and all short-circuit.
has_even = any(x % 2 == 0 for x in [1, 3, 4])
all_positive = all(x > 0 for x in [1, 2])

# Dict and set comprehensions.
squares = {x: x * x for x in range(3)}
letters = {ch for ch in "hello"}

# The walrus operator, for "compute once, test, then reuse".
values = [1, 2, 3]
if (total := sum(values)) > 5:
    leftover = total - 5

# Counting without a loop.
most_common_char, _ = Counter("aabbbcc").most_common(1)[0]

# A grid of the four cardinal moves, used in every BFS and DFS on a grid.
DIRECTIONS = ((1, 0), (-1, 0), (0, 1), (0, -1))

Limits of the language