🎓 A1Learner
🔍
0 of 102 lessons completed

Python itertools: Combinatoric & Grouping

The itertools combinatoric tools list picks and orderings of items.

Pair every choice with product

product pairs each letter with each number. Convert the result to a list before you print it.

from itertools import product

pairs = list(product(["a", "b"], [1, 2]))
print(pairs)

Orderings and picks

permutations care about order, so ab and ba both appear. combinations list each pick once.

from itertools import permutations, combinations

print(list(permutations("ab", 2)))
print(list(combinations("abc", 2)))

Group neighbours with groupby

Sort by the grouping key first. list(group) saves the group before you move on.

from itertools import groupby

words = ["ant", "ape", "bat", "bee"]
words.sort(key=lambda w: w[0])
for letter, group in groupby(words, key=lambda w: w[0]):
    print(letter, list(group))

Running totals with accumulate

accumulate adds as it goes by default. You can pass a function to combine another way.

from itertools import accumulate

print(list(accumulate([1, 2, 3, 4])))
print(list(accumulate([1, 2, 3, 4], lambda a, b: a * b)))

Neighbour pairs with pairwise

pairwise gives each item with the one after it. The last item has no partner.

from itertools import pairwise

print(list(pairwise("abcd")))
print(list(pairwise([10, 20, 30])))

Example

from itertools import groupby, product

def band(score):
    if score >= 90:
        return "A"
    if score >= 80:
        return "B"
    return "C"

scores = [92, 85, 78, 88, 95, 71]
ordered = sorted(scores, key=band)
for name, group in groupby(ordered, key=band):
    print(name, list(group))

print("Dice:")
faces = [1, 2, 3]
for roll in product(faces, repeat=2):
    print(roll)

Key points

  • Combination: a pick of items where order does not matter
  • Permutation: an ordering of items, where order does matter
  • itertools lists these picks and orderings instead of listing choices by hand.
  • Use product(a, b) to pair every item of a with every item of b
  • Use combinations(items, r) to list picks of size r
  • Use groupby to group neighbouring items that share a key
  • Use accumulate to keep a running total as a stream
  • Use pairwise to list each item with the next one
  • groupby only groups neighbours

In this example, band gives each score a grade letter. sorted orders the scores by that letter. groupby then groups neighbouring scores with the same letter. list(group) saves each group before the next group starts. Next, product pairs every dice face with every face. The small dice have faces 1, 2, and 3. So the loop prints nine rolls in total.

Note

Sort by the same key first, or one group can split into two.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Why must you sort before you use groupby?