itertools and functools

Part of the free LogicWiz Python course, Module 12: Iteration and the standard toolkit. The examples and practice run in the app.

Meher wants every possible combo meal of two dishes under ₹50, the running total of today's sales, and the list of orders grouped by dish. Each is one line from itertools.

You'll learn to

Generate arrangements with `permutations`, `combinations` and `product`
Use `accumulate`, `groupby`, `islice` and `chain` on any iterable
Fix some arguments of a function with `partial`
Sort with an old-style compare function using `cmp_to_key`

Learn

Choosing and arranging

From the dishes A, B, C, pick 2:

permutations: order matters. AB and BA are different. 3 × 2 = 6.
combinations: order doesn't matter. AB = BA. 3C2 = 3.
product: one from each list (like nested loops). product("HT", repeat=3) is every result of 3 coin tosses, 2³ = 8.

All three return lazy iterators of tuples, so they are cheap until you loop.

permutations, combinations, product

from itertools import permutations, combinations, product

print(list(permutations("ABC", 2)))
print(list(combinations("ABC", 2)))
print(len(list(product("HT", repeat=3))))

sizes = ["half", "full"]
dishes = ["thali", "biryani"]
for size, dish in product(sizes, dishes):
    print(size, dish)

accumulate, chain, islice

from itertools import accumulate, chain, islice, count

sales = [120, 80, 150, 60]
print(list(accumulate(sales)))          # running total
print(list(accumulate(sales, max)))     # best so far

morning = ["chai", "poha"]
lunch = ["thali", "chai"]
print(list(chain(morning, lunch)))      # one long iterable

print(list(islice(count(101), 5)))      # first 5 of an endless one

islice(it, n) is slicing for iterators and generators, which don't support [:n].

groupby: sort first!

from itertools import groupby

orders = ["chai", "dosa", "chai", "idli", "dosa", "chai"]
for dish, group in groupby(sorted(orders)):
    print(dish, len(list(group)))

print([k for k, g in groupby("aaabccaa")])

groupby only groups items that are next to each other. On unsorted data the same key shows up in several groups. Sort by the same key first: groupby(sorted(data, key=f), key=f). (For plain counting, Counter is simpler.)

functools: partial and cmp_to_key

partial(f, …) makes a new function with some arguments already filled in. partial(bill, gst=5) is a bill that always uses 5% GST. It's a neat fix for the late-binding trap too.

cmp_to_key(compare) turns an old-style compare function, one that returns a negative number, zero or a positive number for a < b, a == b, a > b, into a key= for sorted. You need it when the order depends on both items together and no simple key exists.

partial and cmp_to_key

from functools import partial, cmp_to_key

def bill(amount, gst):
    return round(amount * (1 + gst / 100), 2)

food_bill = partial(bill, gst=5)
print(food_bill(200))
from_binary = partial(int, base=2)
print(from_binary("1011"))

# Largest number from pieces: put a before b if a+b > b+a
def compare(a, b):
    if a + b > b + a:
        return -1
    if a + b < b + a:
        return 1
    return 0

pieces = ["3", "30", "34", "5", "9"]
print("".join(sorted(pieces, key=cmp_to_key(compare))))

Interview corner: "Generators vs lists" (lazy, one pass, low memory), "Counter tricks" (most_common, Counter(a) == Counter(b), c1 - c2), "all subsets / all pairs" (combinations), "largest number" (cmp_to_key) and "run-length encoding" (groupby). These six tools answer most of them.

In one line

permutations (order matters), combinations (it doesn't), product (nested loops), accumulate (running totals), groupby (consecutive only, sort first), islice/chain for iterators; partial pre-fills arguments, cmp_to_key sorts with a compare function.

Interview corner

What is the difference between permutations and combinations in itertools?

permutations(it, r) gives ordered arrangements (AB and BA both appear), n!/(n-r)! of them. combinations(it, r) gives unordered selections (only AB), nCr of them. Use combinations for "all pairs/subsets" and permutations for "all orderings". Both are lazy iterators of tuples.

Why does itertools.groupby sometimes give the same key twice?

groupby groups only consecutive equal keys, like the Unix uniq command. If the data is not sorted by that key, a key that reappears later starts a new group. Sort with the same key first, or use Counter/defaultdict for counting and grouping.

What is functools.partial and when would you use it?

partial(f, *args, **kwargs) returns a new callable with some arguments fixed: int2 = partial(int, base=2). It is handy for callbacks and map (map(partial(round, ndigits=2), prices)), and for avoiding the late-binding closure trap because the values are bound at creation time.

Generators vs lists: give the trade-offs in one breath.

A list stores everything, supports indexing, len() and many passes, but costs memory proportional to its size. A generator stores only its current state, produces items on demand, can be infinite or stop early, but allows a single forward pass with no len() or indexing. Default to a generator for one-pass pipelines over big data.