Modules, packages and the standard library · about 50 min

Standard library tour I: collections, itertools and functools

Why this matters

Python's slogan is "batteries included": the standard library ships with hundreds of modules, so many problems are already solved before you start. Beginners often hand-write 15 lines of dictionary bookkeeping that Counter does in one. This lesson covers three modules that you will use constantly for working with data: collections (specialized containers), itertools (building blocks for iteration) and functools (tools that work on functions). Knowing they exist is half the value; the other half is knowing what they do underneath so you choose the right one.

collections

Counter: counting things

Counting occurrences is so common that it has its own type. A Counter is a dict subclass mapping items to counts:

from collections import Counter

votes = ["red", "blue", "red", "green", "red", "blue"]
c = Counter(votes)
print(c)
print(c["red"], c["purple"])
print(c.most_common(2))

Counter({'red': 3, 'blue': 2, 'green': 1})
3 0
[('red', 3), ('blue', 2)]

Missing items count as 0 instead of raising KeyError. most_common(n) returns the n highest counts; ties keep the order in which items were first seen (because dicts remember insertion order). Counters can be updated and combined:

c.update(["green", "green"])     # add more observations
print(c["green"])                 # 3
print(Counter("hello") + Counter("world"))

3
Counter({'l': 3, 'o': 2, 'h': 1, 'e': 1, 'w': 1, 'r': 1, 'd': 1})

Hand-written, the counting loop would be counts[x] = counts.get(x, 0) + 1 for every item. Counter does the same work, in C, with a clearer name.

defaultdict: dictionaries that fill themselves in

Grouping items means "append to the list for this key, creating the list if needed". With a plain dict that needs a check every time. defaultdict takes a factory, a callable it uses to create a value for any missing key:

from collections import defaultdict

by_initial = defaultdict(list)
for name in ["ada", "alan", "grace", "guido", "linus"]:
    by_initial[name[0]].append(name)

print(dict(by_initial))

{'a': ['ada', 'alan'], 'g': ['grace', 'guido'], 'l': ['linus']}

On by_initial["a"], if "a" is missing, the defaultdict calls list() to make [], stores it under "a", and returns it; .append then modifies that stored list. defaultdict(int) gives 0 for missing keys (a basic counter), defaultdict(set) an empty set.

Warning: merely reading a missing key inserts it. if by_initial["z"]: adds "z": []. Use "z" in by_initial to test membership, and convert to a plain dict before handing it to other code.

deque: a queue with fast ends

A list is a dynamic array: appending or popping at the end is fast, but list.pop(0) or list.insert(0, x) must shift every other element, so it costs time proportional to the length. A deque ("deck", double-ended queue) is a linked sequence of blocks with fast operations at both ends:

from collections import deque

q = deque(["a", "b"])
q.append("c")        # right end
q.appendleft("z")    # left end
print(q.popleft(), q.pop(), list(q))

recent = deque(maxlen=3)        # keeps only the last 3 items
for line in ["l1", "l2", "l3", "l4", "l5"]:
    recent.append(line)
print(list(recent))

z c ['a', 'b']
['l3', 'l4', 'l5']

Use a deque for queues (first in, first out) and for "the last N items", like tail. Indexing into the middle of a deque is slow, so it is not a general list replacement.

namedtuple: tuples with field names

from collections import namedtuple

Point = namedtuple("Point", ["x", "y"])
p = Point(3, 4)
print(p.x, p[1], p)
x, y = p                          # still unpacks like a tuple
print(p._replace(x=10))

3 4 Point(x=3, y=4)
Point(x=10, y=4)

A namedtuple is an immutable tuple whose positions also have names, which makes code that passes records around far more readable than p[0]. (Dataclasses, in the object-oriented module, are the more flexible modern alternative.)

itertools

itertools provides fast building blocks that produce and transform iterators. Most of them are lazy: they produce items one at a time on demand, never building a full list, so they work on huge or even infinite streams.

from itertools import chain, islice, count, product, groupby, accumulate

print(list(chain([1, 2], (3, 4), "ab")))     # one iterable after another
print(list(islice(count(10), 5)))             # first 5 items of an infinite counter
print(list(islice("abcdefg", 2, 6, 2)))       # like slicing [2:6:2], for any iterable
print(list(product("ab", [1, 2])))            # every combination: nested loops
print(list(accumulate([1, 2, 3, 4])))         # running totals

[1, 2, 3, 4, 'a', 'b']
[10, 11, 12, 13, 14]
['c', 'e']
[('a', 1), ('a', 2), ('b', 1), ('b', 2)]
[1, 3, 6, 10]

count(10) would go on forever; islice takes just what you need. That only works because nothing is computed until list() asks for items. product(A, B) replaces two nested for loops; itertools.permutations and itertools.combinations are its cousins.

groupby: grouping consecutive items

groupby(iterable, key=None) walks the input and starts a new group whenever the key changes:

from itertools import groupby

for key, group in groupby("aaabccaa"):
    print(key, list(group))

a ['a', 'a', 'a']
b ['b']
c ['c', 'c']
a ['a', 'a']

Notice a appears twice: groupby only merges adjacent equal items. To group all equal items, sort by the same key first. Each group is an iterator that shares the underlying input with groupby, so consume it before moving to the next group (convert to a list if you need to keep it).

words = ["apple", "bob", "avocado", "banana", "cherry"]
for letter, grp in groupby(sorted(words), key=lambda w: w[0]):
    print(letter, list(grp))

a ['apple', 'avocado']
b ['banana', 'bob']
c ['cherry']

functools

lru_cache: remembering results

Some functions are called again and again with the same arguments. Memoization stores each result the first time and returns the stored value afterwards. functools.lru_cache adds this to any function:

from functools import lru_cache

@lru_cache(maxsize=None)
def fib(n):
    if n < 2:
        return n
    return fib(n - 1) + fib(n - 2)

print(fib(100))
print(fib.cache_info())

354224848179261915075
CacheInfo(hits=98, misses=101, maxsize=None, currsize=101)

Without the cache, fib(100) would make more calls than there are atoms in a mountain; with it, each n is computed once. The @lru_cache(...) line is a decorator: it replaces fib with a wrapper that checks a dictionary keyed by the arguments before calling the original. (You will write decorators yourself in module 13; using them only requires the @ line.) Because the arguments are dictionary keys, they must be hashable: numbers, strings and tuples are; lists and dicts are not. maxsize=128 (the default) keeps only the 128 most recently used results ("LRU" = least recently used is evicted); maxsize=None (or functools.cache) keeps everything.

partial: pre-filling arguments

from functools import partial

def power(base, exponent):
    return base ** exponent

square = partial(power, exponent=2)
print(square(7), list(map(square, [1, 2, 3])))

49 [1, 4, 9]

partial returns a new callable with some arguments fixed, handy for key= and map arguments.

reduce: folding a sequence

reduce(f, items, start) combines items pairwise from the left: f(f(f(start, a), b), c).

from functools import reduce
import operator

print(reduce(operator.mul, [1, 2, 3, 4, 5], 1))   # 5 factorial

For sums, sum() is clearer, and math.prod() exists for products. Reach for reduce only when no named built-in fits.

Common mistakes

Using groupby on unsorted data and getting the same key in several groups. Sort first, by the same key.
Reusing a groupby group after advancing. groups = list(groupby(data)) leaves every group iterator empty, because advancing groupby consumes the shared input. Use [(k, list(g)) for k, g in groupby(data)].
Accidental keys in defaultdicts from lookups, as described above.
list.pop(0) in a loop on a large list: quadratic time. Use deque.popleft().
Caching a function with unhashable arguments (TypeError: unhashable type: 'list') or with side effects: a cached function's body does not run again, so its prints or file writes happen only once.

Recap

Counter counts, defaultdict creates missing values with a factory, deque has fast operations at both ends, namedtuple gives tuple fields names.
itertools functions are lazy iterator building blocks: chain, islice, count, product, accumulate, and groupby for consecutive runs.
lru_cache memoizes by hashable arguments; partial pre-fills arguments; reduce folds a sequence.

Going further: read the "Itertools Recipes" section at the end of the itertools documentation, a collection of short, clever functions built from these pieces. Try help(Counter.subtract) and help(deque.rotate).

Exercises

Exercise 1 · write code

Most common words with Counter

Read all of stdin, lowercase it, split it into words on whitespace, and print the three most common words with their counts, most common first, one per line in the format word: count. Words with equal counts appear in the order they were first seen in the text (this is exactly what Counter.most_common does). If there are fewer than three distinct words, print them all.

For the input the cat and the hat and the bat the output is:

the: 3
and: 2
cat: 1

Use collections.Counter.

File: main.py

Starting point

import sys
from collections import Counter

text = sys.stdin.read()

Example

input

the cat and the hat and the bat

expected output

the: 3
and: 2
cat: 1

Example

input

Red blue
GREEN red Blue green yellow

expected output

red: 2
blue: 2
green: 2

Hint 1

Counter(list_of_words) counts every word in one step.

Hint 2

.most_common(3) returns a list of (word, count) pairs.

Hint 3

Loop with for word, count in ...: and print with an f-string.

Reference solution

import sys
from collections import Counter

words = sys.stdin.read().lower().split()
for word, count in Counter(words).most_common(3):
    print(f"{word}: {count}")

Exercise 2 · write code

Grouping with defaultdict

In solution.py, write a function group_by_length(words) that takes a list of strings and returns a dictionary mapping each word length to the list of words of that length, in the order they appear in the input. Duplicates are kept.

>>> group_by_length(["hi", "cat", "yo", "dog", "a"])
{2: ['hi', 'yo'], 3: ['cat', 'dog'], 1: ['a']}
>>> group_by_length([])
{}

Use collections.defaultdict(list). Returning the defaultdict itself is fine (it compares equal to a plain dict), but converting with dict(...) before returning is good manners: it stops callers from accidentally creating keys just by looking them up.

File: solution.py

Starting point

from collections import defaultdict


def group_by_length(words):
    pass

Hint 1

With groups = defaultdict(list), groups[k].append(x) works even when k is new.

Hint 2

The key is len(word).

Reference solution

from collections import defaultdict


def group_by_length(words):
    """Map each word length to the words of that length, in input order."""
    groups = defaultdict(list)
    for word in words:
        groups[len(word)].append(word)
    return dict(groups)

Exercise 3 · write code

Run-length encoding with groupby

Run-length encoding compresses runs of repeated characters: aaabccdddd becomes a3b1c2d4 (each character of a run followed by the run's length).

Read one line from stdin and print its run-length encoding. Use itertools.groupby, which groups consecutive equal items. An empty line prints an empty line.

Note that runs are consecutive only: aabaa encodes as a2b1a2.

File: main.py

Starting point

from itertools import groupby

line = input()

Example

input

aaabccdddd

expected output

a3b1c2d4

Example

input

aabaa

expected output

a2b1a2

Hint 1

groupby("aaab") yields pairs (key, group_iterator): ('a', <3 a's>), ('b', <1 b>).

Hint 2

The group is an iterator, not a list; len(list(run)) counts it.

Hint 3

Build the pieces in a list and "".join(...) them.

Reference solution

from itertools import groupby

line = input()
parts = [f"{char}{len(list(run))}" for char, run in groupby(line)]
print("".join(parts))

Exercise 4 · quiz

What can lru_cache remember?

from functools import lru_cache

@lru_cache(maxsize=None)
def total(numbers):
    return sum(numbers)

print(total((1, 2, 3)))
print(total([1, 2, 3]))

What happens?

Prints 6 twice
Prints 6, then raises TypeError: unhashable type: 'list'
Raises TypeError on the first call because tuples are not allowed
Prints 6 and then 6 again, but the second call is not cached

Answer

Prints 6, then raises TypeError: unhashable type: 'list'

lru_cache stores results in a dictionary keyed by the arguments, so every argument must be hashable. The tuple (1, 2, 3) is hashable and the call returns 6. A list is mutable and therefore unhashable, so building the cache key fails with TypeError: unhashable type: 'list' before your function body even runs.

Exercise 5 · write code

Memoize an exponential recursion

On a grid, a robot starts in the top-left corner and may only move right or down. The number of distinct paths to the cell at row r, column c (both counted from 0) is:

1 if r == 0 or c == 0 (only one straight line gets there)
otherwise paths(r - 1, c) + paths(r, c - 1)

The starter implements exactly this recursion, but it recomputes the same cells an exponential number of times: paths(16, 16) takes many seconds. Make it fast without changing the recursive logic, by caching results with functools.lru_cache (or functools.cache). The checker computes paths(16, 16) and paths(60, 60) with a 2-second limit.

File: solution.py

Starting point

def paths(r, c):
    if r == 0 or c == 0:
        return 1
    return paths(r - 1, c) + paths(r, c - 1)

Hint 1

Put @lru_cache(maxsize=None) on the line directly above def paths.

Hint 2

Import it first: from functools import lru_cache.

Reference solution

from functools import lru_cache


@lru_cache(maxsize=None)
def paths(r, c):
    """Number of right/down paths from (0, 0) to (r, c)."""
    if r == 0 or c == 0:
        return 1
    return paths(r - 1, c) + paths(r, c - 1)