|
| 1 | +""" |
| 2 | +Itertools provides a collection of tools for handling iterators. This |
| 3 | +module demonstrates how to use itertools functions to efficiently work |
| 4 | +with sequences, combine iterables, and create infinite iterators. |
| 5 | +""" |
| 6 | + |
| 7 | +import itertools |
| 8 | + |
| 9 | + |
| 10 | +def main() -> None: |
| 11 | + # chain() combines multiple iterables into a single iterator |
| 12 | + letters = ["a", "b", "c"] |
| 13 | + numbers = [1, 2, 3] |
| 14 | + combined = list(itertools.chain(letters, numbers)) |
| 15 | + assert combined == ["a", "b", "c", 1, 2, 3] |
| 16 | + |
| 17 | + # cycle() creates an infinite iterator that cycles through the elements |
| 18 | + # We'll use islice to take only the first few elements |
| 19 | + cycled = list(itertools.islice(itertools.cycle(["A", "B"]), 6)) |
| 20 | + assert cycled == ["A", "B", "A", "B", "A", "B"] |
| 21 | + |
| 22 | + # repeat() creates an iterator that repeats a value infinitely |
| 23 | + repeated = list(itertools.islice(itertools.repeat("hello"), 4)) |
| 24 | + assert repeated == ["hello", "hello", "hello", "hello"] |
| 25 | + |
| 26 | + # count() creates an infinite iterator that counts up from a start value |
| 27 | + counted = list(itertools.islice(itertools.count(10), 5)) |
| 28 | + assert counted == [10, 11, 12, 13, 14] |
| 29 | + |
| 30 | + # islice() allows slicing of iterators (similar to list slicing) |
| 31 | + data = itertools.count() # infinite count from 0 |
| 32 | + sliced = list(itertools.islice(data, 2, 8, 2)) # start=2, stop=8, step=2 |
| 33 | + assert sliced == [2, 4, 6] |
| 34 | + |
| 35 | + # tee() creates multiple independent iterators from one |
| 36 | + original = iter([1, 2, 3, 4, 5]) |
| 37 | + iter1, iter2 = itertools.tee(original, 2) |
| 38 | + list1 = list(iter1) |
| 39 | + list2 = list(iter2) |
| 40 | + assert list1 == [1, 2, 3, 4, 5] |
| 41 | + assert list2 == [1, 2, 3, 4, 5] |
| 42 | + |
| 43 | + # groupby() groups consecutive equal elements |
| 44 | + group_data = [1, 1, 2, 2, 2, 3, 1, 1] |
| 45 | + groups = [(key, list(group)) for key, group in itertools.groupby(group_data)] |
| 46 | + assert groups == [(1, [1, 1]), (2, [2, 2, 2]), (3, [3]), (1, [1, 1])] |
| 47 | + |
| 48 | + # product() creates cartesian product of input iterables |
| 49 | + colors = ["red", "blue"] |
| 50 | + sizes = ["S", "M"] |
| 51 | + combinations = list(itertools.product(colors, sizes)) |
| 52 | + assert combinations == [("red", "S"), ("red", "M"), ("blue", "S"), ("blue", "M")] |
| 53 | + |
| 54 | + # permutations() generates all possible orderings |
| 55 | + perms = list(itertools.permutations([1, 2, 3], 2)) # length 2 permutations |
| 56 | + assert len(perms) == 6 # 3! / (3-2)! = 6 |
| 57 | + assert (1, 2) in perms |
| 58 | + assert (2, 1) in perms |
| 59 | + |
| 60 | + # combinations() generates combinations (order doesn't matter) |
| 61 | + combos = list(itertools.combinations([1, 2, 3, 4], 2)) |
| 62 | + assert len(combos) == 6 # C(4,2) = 6 |
| 63 | + assert (1, 2) in combos |
| 64 | + assert (2, 1) not in combos # order doesn't matter |
| 65 | + |
| 66 | + |
| 67 | +if __name__ == "__main__": |
| 68 | + main() |
0 commit comments