Python · Syntax · Intermediate

itertools

10 tasks

Master Python's itertools module: tools for working with iterators, generating sequences, and solving combinatorics problems efficiently.

Quick topic start and explanations before exercises (exercises below):

itertools: Infinite Iterators & Slicing

#
The `itertools` module provides memory-efficient tools for working with sequences and iterators. All functions return **iterators** — they generate values on demand instead of building full lists in memory. ```python import itertools ``` ## Chaining iterables ```python # chain: treat multiple iterables as one continuous sequence result = list(itertools.chain([1, 2], [3, 4], [5])) print(result) # [1, 2, 3, 4, 5] # chain.from_iterable: flatten one level nested = [[1, 2], [3, 4], [5, 6]] flat = list(itertools.chain.from_iterable(nested)) print(flat) # [1, 2, 3, 4, 5, 6] ``` ## islice: slicing any iterator `islice` lets you slice iterators that do not support `[start:stop]`: ```python # islice(iterable, stop) # islice(iterable, start, stop) # islice(iterable, start, stop, step) data = range(100) first_five = list(itertools.islice(data, 5)) print(first_five) # [0, 1, 2, 3, 4] skip_then_take = list(itertools.islice(data, 10, 15)) print(skip_then_take) # [10, 11, 12, 13, 14] ``` ## Infinite iterators These never stop — always use with `islice` or a `break` condition: ```python # count(start=0, step=1): counts forever counter = itertools.count(10, 2) print(list(itertools.islice(counter, 5))) # [10, 12, 14, 16, 18] # cycle(iterable): repeats the sequence indefinitely colors = itertools.cycle(["red", "green", "blue"]) print(list(itertools.islice(colors, 7))) # ["red", "green", "blue", "red", "green", "blue", "red"] # repeat(object, times=None): repeat one value fives = itertools.repeat(5, 4) print(list(fives)) # [5, 5, 5, 5] # repeat without times is infinite — useful with map/zip result = list(map(lambda x, n: x ** n, [1, 2, 3, 4], itertools.repeat(2))) print(result) # [1, 4, 9, 16] ``` ## zip_longest: zip with unequal lengths ```python a = [1, 2, 3] b = ["a", "b"] # Built-in zip stops at the shorter sequence: print(list(zip(a, b))) # [(1, "a"), (2, "b")] # zip_longest fills missing values: print(list(itertools.zip_longest(a, b, fillvalue=0))) # [(1, "a"), (2, "b"), (3, 0)] ``` ## accumulate: running totals ```python import operator data = [1, 2, 3, 4, 5] # Running sum (default) print(list(itertools.accumulate(data))) # [1, 3, 6, 10, 15] # Running product print(list(itertools.accumulate(data, operator.mul))) # [1, 2, 6, 24, 120] # Running maximum print(list(itertools.accumulate([3, 1, 4, 1, 5, 9, 2, 6], max))) # [3, 3, 4, 4, 5, 9, 9, 9] ``` ## takewhile and dropwhile ```python data = [1, 3, 5, 2, 4, 6] # takewhile: take elements while condition is True, then stop print(list(itertools.takewhile(lambda x: x < 4, data))) # [1, 3] — stops at 5 because 5 >= 4? wait, 5 is not < 4, stops at 5 # Corrected: print(list(itertools.takewhile(lambda x: x % 2 != 0, data))) # [1, 3, 5] — stops at 2 (first even number) # dropwhile: skip elements while condition is True, then yield rest print(list(itertools.dropwhile(lambda x: x % 2 != 0, data))) # [2, 4, 6] — skips odd numbers at start, yields from first even ```

Combinatorics with itertools

#
## product: Cartesian product ```python import itertools # All combinations of sizes and colors sizes = ["S", "M", "L"] colors = ["red", "blue"] variants = list(itertools.product(sizes, colors)) print(variants) # [("S","red"), ("S","blue"), ("M","red"), ("M","blue"), ("L","red"), ("L","blue")] # product with repeat: rolling a 2-dice dice = list(itertools.product(range(1, 7), repeat=2)) print(len(dice)) # 36 print(dice[:3]) # [(1, 1), (1, 2), (1, 3)] ``` ## permutations: ordered arrangements ```python # permutations(iterable, r=None): all ordered r-length arrangements letters = ["A", "B", "C"] # All 3-length permutations (r defaults to len) perms = list(itertools.permutations(letters)) print(len(perms)) # 6 (3! = 6) print(perms) # [("A","B","C"), ("A","C","B"), ("B","A","C"), ...] # 2-length permutations perms2 = list(itertools.permutations(letters, 2)) print(perms2) # [("A","B"), ("A","C"), ("B","A"), ("B","C"), ("C","A"), ("C","B")] ``` ## combinations: unordered selections ```python # combinations(iterable, r): r-length, no repeats, order does not matter people = ["Alice", "Bob", "Carol", "David"] pairs = list(itertools.combinations(people, 2)) print(len(pairs)) # 6 (C(4,2) = 6) print(pairs[:3]) # [("Alice","Bob"), ("Alice","Carol"), ("Alice","David")] # combinations_with_replacement: allows repeated elements # Example: all pairs of dice (with ties allowed) dice_pairs = list(itertools.combinations_with_replacement(range(1, 4), 2)) print(dice_pairs) # [(1,1), (1,2), (1,3), (2,2), (2,3), (3,3)] ``` ## groupby: group consecutive elements `groupby` groups **consecutive** elements with the same key. The input must be sorted by the key first. ```python data = [ {"name": "Alice", "dept": "Engineering"}, {"name": "Bob", "dept": "Engineering"}, {"name": "Carol", "dept": "Marketing"}, {"name": "David", "dept": "Marketing"}, {"name": "Eve", "dept": "Engineering"}, ] # Sort first by the grouping key data.sort(key=lambda x: x["dept"]) for dept, members in itertools.groupby(data, key=lambda x: x["dept"]): names = [m["name"] for m in members] print(f"{dept}: {names}") # Engineering: ["Alice", "Bob", "Eve"] # Marketing: ["Carol", "David"] ``` ## filterfalse: opposite of filter ```python nums = [1, 2, 3, 4, 5, 6, 7, 8] # filter: keep elements where condition is True evens = list(filter(lambda x: x % 2 == 0, nums)) print(evens) # [2, 4, 6, 8] # filterfalse: keep elements where condition is False odds = list(itertools.filterfalse(lambda x: x % 2 == 0, nums)) print(odds) # [1, 3, 5, 7] ```

itertools Reference

#
## Infinite iterators | Function | Signature | What it does | |---|---|---| | `count` | `count(start=0, step=1)` | Count from start indefinitely | | `cycle` | `cycle(iterable)` | Repeat iterable indefinitely | | `repeat` | `repeat(obj, times=None)` | Repeat obj, optionally N times | ## Slicing iterators | Function | Signature | What it does | |---|---|---| | `islice` | `islice(it, stop)` or `islice(it, start, stop[, step])` | Slice any iterator | | `takewhile` | `takewhile(pred, it)` | Yield while pred is True | | `dropwhile` | `dropwhile(pred, it)` | Skip while pred is True, then yield rest | | `filterfalse` | `filterfalse(pred, it)` | Yield where pred is False | ## Combining iterators | Function | Signature | What it does | |---|---|---| | `chain` | `chain(*iterables)` | Concatenate iterables | | `chain.from_iterable` | `chain.from_iterable(iterable)` | Flatten one level | | `zip_longest` | `zip_longest(*its, fillvalue=None)` | Zip, padding shorter sequences | | `accumulate` | `accumulate(it, func=operator.add)` | Running total with func | | `groupby` | `groupby(it, key=None)` | Group consecutive equal-key elements | ## Combinatorics | Function | Signature | Count | |---|---|---| | `product` | `product(*its, repeat=1)` | n1 × n2 × ... | | `permutations` | `permutations(it, r=None)` | n! / (n-r)! | | `combinations` | `combinations(it, r)` | n! / (r! × (n-r)!) | | `combinations_with_replacement` | `combinations_with_replacement(it, r)` | (n+r-1)! / (r! × (n-1)!) | ## groupby gotcha ```python # WRONG: groupby only groups CONSECUTIVE equal elements data = [1, 2, 1, 2] for k, g in itertools.groupby(data): print(k, list(g)) # 1 [1] # 2 [2] # 1 [1] ← NOT grouped with the first 1! # 2 [2] # CORRECT: sort first data.sort() for k, g in itertools.groupby(data): print(k, list(g)) # 1 [1, 1] # 2 [2, 2] ``` ## Common patterns ```python import itertools, operator # Flatten a nested list flat = list(itertools.chain.from_iterable(nested)) # Running sum running = list(itertools.accumulate(values)) # First N items from generator head = list(itertools.islice(gen, N)) # All pairs (without self-pairs) pairs = list(itertools.combinations(items, 2)) # Infinite counter with zip numbered = list(itertools.islice( zip(itertools.count(1), items), len(items) )) ```

Exercises:

01

Flatten Nested List

#
Write a function `flatten(nested)` that takes a list of lists and returns a flat list using `itertools.chain.from_iterable`. Example: `flatten([[1, 2], [3], [4, 5, 6]])` → `[1, 2, 3, 4, 5, 6]`.
from itertools import chain


def flatten(nested):
    pass


print(flatten([[1, 2], [3], [4, 5, 6]]))
# [1, 2, 3, 4, 5, 6]
Solution 
from itertools import chain


def flatten(nested):
    return list(chain.from_iterable(nested))


print(flatten([[1, 2], [3], [4, 5, 6]]))
# [1, 2, 3, 4, 5, 6]
02

First N from Generator

#
Write a function `first_n(gen, n)` that returns the first `n` elements from any generator or iterator as a list, using `itertools.islice`. Example: `first_n((x**2 for x in range(100)), 5)` → `[0, 1, 4, 9, 16]`.
from itertools import islice


def first_n(gen, n):
    pass


print(first_n((x**2 for x in range(100)), 5))
# [0, 1, 4, 9, 16]
Solution 
from itertools import islice


def first_n(gen, n):
    return list(islice(gen, n))


print(first_n((x**2 for x in range(100)), 5))
# [0, 1, 4, 9, 16]
03

Round-Robin Scheduler

#
Write a function `round_robin(tasks)` that takes a list of task names and returns an infinite cycle iterator over them using `itertools.cycle`. Then use `itertools.islice` to extract the first 8 assignments. Example: `list(islice(round_robin(["A", "B", "C"]), 8))` → `["A", "B", "C", "A", "B", "C", "A", "B"]`.
from itertools import cycle, islice


def round_robin(tasks):
    pass


print(list(islice(round_robin(["A", "B", "C"]), 8)))
# ["A", "B", "C", "A", "B", "C", "A", "B"]
Solution 
from itertools import cycle, islice


def round_robin(tasks):
    return cycle(tasks)


print(list(islice(round_robin(["A", "B", "C"]), 8)))
# ["A", "B", "C", "A", "B", "C", "A", "B"]
04

Merge Unequal Sequences

#
Write a function `merge_sequences(seq1, seq2, seq3, fill=None)` that combines three sequences element by element using `itertools.zip_longest`, with `fill` as the missing value. Return the result as a list of tuples. Example: `merge_sequences([1,2,3], ["a","b"], [True], fill=0)` → `[(1,"a",True), (2,"b",0), (3,0,0)]`.
from itertools import zip_longest


def merge_sequences(seq1, seq2, seq3, fill=None):
    pass


print(merge_sequences([1, 2, 3], ["a", "b"], [True], fill=0))
# [(1, "a", True), (2, "b", 0), (3, 0, 0)]
Solution 
from itertools import zip_longest


def merge_sequences(seq1, seq2, seq3, fill=None):
    return list(zip_longest(seq1, seq2, seq3, fillvalue=fill))


print(merge_sequences([1, 2, 3], ["a", "b"], [True], fill=0))
# [(1, "a", True), (2, "b", 0), (3, 0, 0)]
05

Running Maximum

#
Write a function `running_max(numbers)` that returns a list where each element is the maximum value seen so far in `numbers`. Use `itertools.accumulate`. Example: `running_max([3, 1, 4, 1, 5, 9, 2, 6])` → `[3, 3, 4, 4, 5, 9, 9, 9]`.
from itertools import accumulate
import operator


def running_max(numbers):
    pass


print(running_max([3, 1, 4, 1, 5, 9, 2, 6]))
# [3, 3, 4, 4, 5, 9, 9, 9]
Solution 
from itertools import accumulate
import operator


def running_max(numbers):
    return list(accumulate(numbers, func=max))


print(running_max([3, 1, 4, 1, 5, 9, 2, 6]))
# [3, 3, 4, 4, 5, 9, 9, 9]
06

Cartesian Product Grid

#
Write a function `grid(rows, cols)` that returns all (row, col) coordinate pairs for a grid of given dimensions, using `itertools.product`. Example: `grid(2, 3)` → `[(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]`.
from itertools import product


def grid(rows, cols):
    pass


print(grid(2, 3))
# [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]
Solution 
from itertools import product


def grid(rows, cols):
    return list(product(range(rows), range(cols)))


print(grid(2, 3))
# [(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)]
07

All Pairs (Combinations)

#
Write a function `all_pairs(items)` that returns all unique unordered pairs of items (no item paired with itself). Use `itertools.combinations`. Example: `all_pairs(["A", "B", "C", "D"])` → `[("A","B"), ("A","C"), ("A","D"), ("B","C"), ("B","D"), ("C","D")]`.
from itertools import combinations


def all_pairs(items):
    pass


print(all_pairs(["A", "B", "C", "D"]))
# [("A","B"), ("A","C"), ("A","D"), ("B","C"), ("B","D"), ("C","D")]
Solution 
from itertools import combinations


def all_pairs(items):
    return list(combinations(items, 2))


print(all_pairs(["A", "B", "C", "D"]))
# [("A", "B"), ("A", "C"), ("A", "D"), ("B", "C"), ("B", "D"), ("C", "D")]
08

Group by Category

#
Write a function `group_by(items, key_func)` that groups a list of items by a key function and returns a dict mapping each key to a list of items. Use `itertools.groupby` (remember to sort first). Example: with `items = ["apple", "ant", "banana", "bear", "cherry"]` and `key_func = lambda w: w[0]`, result is `{"a": ["ant", "apple"], "b": ["banana", "bear"], "c": ["cherry"]}`.
from itertools import groupby


def group_by(items, key_func):
    pass


words = ["apple", "ant", "banana", "bear", "cherry"]
print(group_by(words, lambda w: w[0]))
# {"a": ["ant", "apple"], "b": ["banana", "bear"], "c": ["cherry"]}
Solution 
from itertools import groupby


def group_by(items, key_func):
    sorted_items = sorted(items, key=key_func)
    result = {}
    for key, group in groupby(sorted_items, key=key_func):
        result[key] = list(group)
    return result


words = ["apple", "ant", "banana", "bear", "cherry"]
print(group_by(words, lambda w: w[0]))
# {"a": ["ant", "apple"], "b": ["banana", "bear"], "c": ["cherry"]}
09

takewhile: Read Until Sentinel

#
Write a function `read_until_empty(lines)` that takes a list of strings and returns only the lines before the first empty string (or line containing only whitespace), using `itertools.takewhile`. Example: `read_until_empty(["hello", "world", "", "more", "text"])` → `["hello", "world"]`.
from itertools import takewhile


def read_until_empty(lines):
    pass


print(read_until_empty(["hello", "world", "", "more", "text"]))
# ["hello", "world"]
Solution 
from itertools import takewhile


def read_until_empty(lines):
    return list(takewhile(lambda line: line.strip() != "", lines))


print(read_until_empty(["hello", "world", "", "more", "text"]))
# ["hello", "world"]
10

Password Generator

#
Write a function `generate_passwords(chars, length, count)` that generates `count` unique random passwords of `length` characters, where each password is built from `itertools.product(chars, repeat=length)`. Return the first `count` products as joined strings. Example: `generate_passwords("ab", 2, 4)` → `["aa", "ab", "ba", "bb"]` (use product in order, no randomness needed).
from itertools import product, islice


def generate_passwords(chars, length, count):
    pass


print(generate_passwords("ab", 2, 4))
# ["aa", "ab", "ba", "bb"]
Solution 
from itertools import product, islice


def generate_passwords(chars, length, count):
    return ["".join(p) for p in islice(product(chars, repeat=length), count)]


print(generate_passwords("ab", 2, 4))
# ["aa", "ab", "ba", "bb"]