111. Itertools Module

🔗 A data processing pipeline uses itertools for efficient iteration. The developer forgot to wrap chain() in list(), and called groupby() on unsorted data — groupby only groups consecutive identical keys!

💡 Fun fact: The itertools module was added in Python 2.3 (2003) and is modeled after APL, Haskell, and SML functional programming idioms. All itertools functions return lazy iterators — they produce values one at a time and never build the entire result in memory. The Python docs include a section called “Itertools Recipes” with dozens of useful combinations like sliding_window, pairwise, and batched.

⚠️ Watch out: groupby() groups ONLY CONSECUTIVE identical keys. If your data is ["Alice", "Bob", "Anna"], groupby sees three separate groups (A, B, A) — not two. You MUST sort the data by the same key function before calling groupby. This is the single most common itertools mistake.

🤔 Think about it: chain(*nested) unpacks and flattens a list of lists lazily. If nested has 1,000 lists each with 1,000 items, how much memory does chain(*nested) use compared to [item for lst in nested for item in lst]? When would the list comprehension be preferable despite using more memory?

Learning objectives

Use chain() to flatten nested iterables into a single sequence
Understand that itertools returns lazy iterators — wrap in list() when needed
Sort data before groupby() to ensure all matching keys are grouped together
Use islice() to efficiently take the first N items from any iterable
Use product() to generate all combinations from multiple sequences

Key concepts

chain(*iterables) — concatenate iterables lazily
islice(iterable, n) — take first n items
groupby(sorted_data, key) — group consecutive equal keys
product(*iterables) — cartesian product
Iterator vs list — itertools is lazy, must call list() to materialize

Try it

from itertools import chain, islice, groupby
def flatten_lists(nested):
"""Flatten a list of lists into a single list."""
# 💡 : hint : -010 xp chain(*nested) returns a lazy iterator — wrap it in list() to get an indexable list
return chain(*nested) # BUG 1: chain(*nested) returns an iterator, not a list
def first_n(iterable, n):
"""Return the first n items from any iterable as a list."""
return list(islice(iterable, n))
def group_by_first_letter(names):
"""Group names by their first letter. Returns dict of letter -> [names]."""
# 💡 : hint : -010 xp groupby only groups CONSECUTIVE identical keys — sort by the same key first so all matching names are adjacent
groups = {}
for letter, members in groupby(names, key=lambda n: n[0].upper()): # BUG 2: unsorted — same-letter names in different positions form separate groups
groups[letter] = list(members)
return groups
def cartesian_pairs(list_a, list_b):
"""Return all (a, b) pairs from list_a x list_b."""
from itertools import product
return list(product(list_a, list_b))

Concept detail

itertools — Efficient Iteration Tools

The itertools module provides fast, memory-efficient iteration building blocks.

chain — Flatten iterables

from itertools import chain
flat = list(chain([1, 2], [3, 4], [5]))  # [1, 2, 3, 4, 5]
# chain(*nested) unpacks a list of lists
flat = list(chain(*[[1, 2], [3, 4]]))    # [1, 2, 3, 4]

Key point: chain() returns an iterator. Always wrap in list() if you need indexing or len().

islice — Take first N items

from itertools import islice
first5 = list(islice(range(1000000), 5))  # [0, 1, 2, 3, 4]

Works on any iterable without creating the full sequence in memory.

groupby — Group consecutive items

from itertools import groupby
data = sorted(["Alice", "Bob", "Anna"], key=lambda x: x[0])
for key, group in groupby(data, key=lambda x: x[0]):
    print(key, list(group))
# A ['Alice', 'Anna']
# B ['Bob']

Critical: groupby only groups consecutive identical keys. Always sort first.

product — Cartesian product

from itertools import product
pairs = list(product([1, 2], ["a", "b"]))
# [(1, 'a'), (1, 'b'), (2, 'a'), (2, 'b')]

count — Infinite counter

from itertools import count
counter = count(start=1, step=2)  # 1, 3, 5, 7, ...
first10 = list(islice(counter, 10))

Solution

from itertools import chain, islice, groupby, product

def flatten_lists(nested):
    """Flatten a list of lists into a single list."""
    # FIX 1: wrap in list() so callers get an indexable list
    return list(chain(*nested))

def first_n(iterable, n):
    """Return the first n items from any iterable as a list."""
    return list(islice(iterable, n))

def group_by_first_letter(names):
    """Group names by their first letter. Returns dict of letter -> [names]."""
    # FIX 2: sort first so groupby sees consecutive equal keys
    sorted_names = sorted(names, key=lambda n: n[0].upper())
    groups = {}
    for letter, members in groupby(sorted_names, key=lambda n: n[0].upper()):
        groups[letter] = list(members)
    return groups

def cartesian_pairs(list_a, list_b):
    """Return all (a, b) pairs from list_a x list_b."""
    return list(product(list_a, list_b))

Tests

def test_flatten_lists_returns_list():
    result = flatten_lists([[1, 2], [3, 4]])
    assert isinstance(result, list), "flatten_lists must return a list, not an iterator"

def test_flatten_lists_correct_elements():
    result = flatten_lists([[1, 2], [3, 4], [5]])
    assert result == [1, 2, 3, 4, 5]

def test_flatten_lists_empty():
    assert flatten_lists([]) == []

def test_first_n_returns_correct_count():
    assert first_n(range(100), 5) == [0, 1, 2, 3, 4]

def test_first_n_shorter_than_n():
    assert first_n([1, 2], 10) == [1, 2]

def test_group_by_first_letter_all_groups_present():
    names = ["Alice", "Bob", "Anna", "Charlie", "Beth"]
    result = group_by_first_letter(names)
    # Without sorting, "Anna" and "Alice" would be in separate groups
    assert "A" in result
    assert "B" in result
    assert "C" in result
    assert set(result["A"]) == {"Alice", "Anna"}

def test_group_by_first_letter_no_duplicate_keys():
    names = ["Alice", "Anna", "Bob", "Beth", "Alice"]
    result = group_by_first_letter(names)
    # Each letter key should appear exactly once
    assert len([k for k in result if k == "A"]) == 1
    assert len([k for k in result if k == "B"]) == 1

def test_cartesian_pairs():
    result = cartesian_pairs([1, 2], ["a", "b"])
    assert len(result) == 4
    assert (1, "a") in result
    assert (2, "b") in result