Combinations in Python

Ben Cook • Posted 2020-12-28 • Last updated 2021-10-15

If you want to see how to create combinations without itertools in Python, jump to this section.

Combinations

A combination is a selection of elements from a set such that order doesn’t matter. Say we have a list [1, 2, 3], the 2-combinations of this set are [(1, 2), (1, 3), (2, 3)]. Notice that order doesn’t matter. Once we have (1, 2) in the set, we don’t also get (2, 1). By default, combinations are typically defined to be without replacement. This means that we’ll never see (1, 1) – once the 1 has been drawn it is not replaced.

You can also have combinations with replacement. The 2-combinations (with replacement) of the list [1, 2, 3] are [(1, 1), (1, 2), (1, 3), (2, 2), (2, 3), (3, 3)]. In this case, numbers are replaced after they’re drawn.

There’s one important note before we jump into implementations of this operation in Python. The combinations API from itertools treats list index as the element being drawn. This means any iterable can be treated like a set (since all indices are unique). But it’s important to realize that if you pass in [1, 1, 2], the elements will not be de-duped for you. The 2-combinations of [1, 1, 2] according to the itertools combinations API is [(1, 1), (1, 2), (1, 2)].

Approaches

Combinations in itertools

It’s extremely easy to generate combinations in Python with itertools. The following generates all 2-combinations of the list [1, 2, 3]:

import itertools

sequence = [1, 2, 3]
itertools.combinations(sequence, 2)

# Expected result
# <itertools.combinations at 0x7fcbd25cc3b8>

The combinations() function returns an iterator. This is what you want if you plan to loop through the combinations. But you can convert it into a list if you want all the combinations in memory:

list(itertools.combinations(sequence, 2))

# Expected result
# [(1, 2), (1, 3), (2, 3)]

A useful property of the combinations() function is that it takes any iterable as the first argument. This means you can pass lazy sequences[1] in:

list(itertools.combinations(range(3), 2))

# Expected result
# [(0, 1), (0, 2), (1, 2)]

Combinations with replacement in itertools

It’s also very easy to generate combinations with replacement:

list(itertools.combinations_with_replacement(sequence, 2))

# Expected result
# [(1, 1), (1, 2), (1, 3), (2, 2), (2, 3), (3, 3)]

The interface for combinations_with_replacement() is the same as combinations().

Combinations without itertools

Once in a while, you might want to generate combinations without using itertools. Maybe you want to change the API slightly — say, returning a list instead of an iterator, or you might want to operate on a NumPy array.

Under the hood, Python uses a C implementation of the combinations algorithm. But the documentation provides a helpful Python implementation you can use, reproduced here for convenience:

def combinations(iterable, r):
    # combinations('ABCD', 2) --> AB AC AD BC BD CD
    # combinations(range(4), 3) --> 012 013 023 123
    pool = tuple(iterable)
    n = len(pool)
    if r > n:
        return
    indices = list(range(r))
    yield tuple(pool[i] for i in indices)
    while True:
        for i in reversed(range(r)):
            if indices[i] != i + n - r:
                break
        else:
            return
        indices[i] += 1
        for j in range(i+1, r):
            indices[j] = indices[j-1] + 1
        yield tuple(pool[i] for i in indices)

Combinations without replacement (and without itertools)

The Python docs also give us a Python-only implementation of combinations_with_replacement():

def combinations_with_replacement(iterable, r):
    # combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC
    pool = tuple(iterable)
    n = len(pool)
    if not n and r:
        return
    indices = [0] * r
    yield tuple(pool[i] for i in indices)
    while True:
        for i in reversed(range(r)):
            if indices[i] != n - 1:
                break
        else:
            return
        indices[i:] = [indices[i] + 1] * (r - i)
        yield tuple(pool[i] for i in indices)

Notes

[1]: Technically, range() does not return an iterator.