Python Itertools Combinations

Source

What?

This is a brief note on the python itertools.combinations function.

from itertools import combinations

To make sure it's behaving the way I think it will I'll also use some other built-ins.

from itertools import count
from math import factorial
from math import comb as n_choose_k

Okay, but what?

The combinations function takes two arguments:

  • iterable: the source that you want to make combinations from
  • r: the length of the sequences you want to make

It returns all the r-length subsequences in the iterable (combinations).

A Count Function

The number of combinations of length k for a sequence of length n is the Binomial Coefficient.

\begin{align} \binom{n}{k} &= \frac{n(n-1) \cdots (n - k + 1)}{k(k-1) \cdots 1}\\ &= \frac{n!}{k!(n-k)!} \end{align}

I originally implemented a function for this but it turns out it got added to python as of 3.8 (number-theoretic functions). Here's my translation for reference. The calculation is actually really fast so you don't gain a lot from these toy problems using the standard library, but since it's there, why not?

def n_choose_k_one(n: int, k: int) -> int:
    """The Binomial Coefficient

    Args:
     n: the length of the source set
     k: the length of the combinations

    Returns:
     the number of combinations in the sequence
    """
    return factorial(n) / (factorial(k) * factorial(n - k))

Example One

This is given as an example in the documentation.

EXPECTEDS = "AB AC AD BC BD CD".split()
SOURCE = "ABCD"
k = 2
for combination, expected, total in zip(combinations(iterable=SOURCE, r=k), EXPECTEDS, count(start=1)):
    print(f" - {combination}")
    assert "".join(combination) == expected

n = len(SOURCE)
assert total == n_choose_k(n, k)

- ('A', 'B')
- ('A', 'C')
- ('A', 'D')
- ('B', 'C')
- ('B', 'D')
- ('C', 'D')

Note that it treated the string as a sequence and returns a tuple, not a sub-string. Also note that it doesn't swap orders, so "D" is never a first entry, for instance.

Example Two

This one is also from the documentation.

SOURCE = list(range(4))
n = len(SOURCE)
EXPECTEDS = [(0, 1, 2),
             (0, 1, 3),
             (0, 2, 3),
             (1, 2, 3)]
k = 3

for combination, expected, total in zip(combinations(SOURCE, k), EXPECTEDS, count(start=1)):
    print(combination)
    assert combination == expected

assert total == n_choose_k(n, k)

(0, 1, 2)
(0, 1, 3)
(0, 2, 3)
(1, 2, 3)