itertools.py

"""
The `itertools` module provides a suite of fast, memory-efficient tools
for constructing and manipulating iterators.

By utilizing lazy evaluation, `itertools` allows programs to process large or
even infinite data streams without loading the entire dataset into memory. Its
functions are written in C, making them faster than equivalent pure Python loops.
These tools are generally categorized into infinite iterators (like `count`, `cycle`),
terminating iterators (like `accumulate`, `chain`), and combinatoric iterators
(like `product`, `permutations`).
"""

import itertools


def main() -> None:
    # chain() combines multiple iterables into a single iterator
    letters = ["a", "b", "c"]
    numbers = [1, 2, 3]
    combined = list(itertools.chain(letters, numbers))
    assert combined == ["a", "b", "c", 1, 2, 3]

    # cycle() creates an infinite iterator that cycles through the elements
    # We'll use islice to take only the first few elements
    cycled = list(itertools.islice(itertools.cycle(["A", "B"]), 6))
    assert cycled == ["A", "B", "A", "B", "A", "B"]

    # repeat() creates an iterator that repeats a value infinitely
    repeated = list(itertools.islice(itertools.repeat("hello"), 4))
    assert repeated == ["hello", "hello", "hello", "hello"]

    # count() creates an infinite iterator that counts up from a start value
    counted = list(itertools.islice(itertools.count(10), 5))
    assert counted == [10, 11, 12, 13, 14]

    # islice() allows slicing of iterators (similar to list slicing)
    data = itertools.count()  # infinite count from 0
    sliced = list(itertools.islice(data, 2, 8, 2))  # start=2, stop=8, step=2
    assert sliced == [2, 4, 6]

    # tee() creates multiple independent iterators from one
    original = iter([1, 2, 3, 4, 5])
    iter1, iter2 = itertools.tee(original, 2)
    list1 = list(iter1)
    list2 = list(iter2)
    assert list1 == [1, 2, 3, 4, 5]
    assert list2 == [1, 2, 3, 4, 5]

    # groupby() groups consecutive equal elements
    group_data = [1, 1, 2, 2, 2, 3, 1, 1]
    groups = [(key, list(group)) for key, group in itertools.groupby(group_data)]
    assert groups == [(1, [1, 1]), (2, [2, 2, 2]), (3, [3]), (1, [1, 1])]

    # product() creates cartesian product of input iterables
    colors = ["red", "blue"]
    sizes = ["S", "M"]
    combinations = list(itertools.product(colors, sizes))
    assert combinations == [("red", "S"), ("red", "M"), ("blue", "S"), ("blue", "M")]

    # permutations() generates all possible orderings
    perms = list(itertools.permutations([1, 2, 3], 2))  # length 2 permutations
    assert len(perms) == 6  # 3! / (3-2)! = 6
    assert (1, 2) in perms
    assert (2, 1) in perms

    # combinations() generates combinations (order doesn't matter)
    combos = list(itertools.combinations([1, 2, 3, 4], 2))
    assert len(combos) == 6  # C(4,2) = 6
    assert (1, 2) in combos
    assert (2, 1) not in combos  # order doesn't matter


if __name__ == "__main__":
    main()