Python itertools模块详解

Python itertools模块提供了一组功能强大、效率高的工具，用于处理各种迭代器（iterators）。本文将详细讲解 itertools 模块中常用的函数及其用法。

itertools.count

itertools.count(start=0, step=1) 函数生成一个无限序列，从 start 开始，步长为 step。start 和 step 都是可选参数，默认值分别为 0 和 1。

import itertools

for i in itertools.count(10, 2):
    if i > 20:
        break
    else:
        print(i)

输出：

10
12
14
16
18
20

itertools.cycle

itertools.cycle(iterable) 函数生成一个无限序列，从 iterable 中不断取元素循环输出。

import itertools

lst = ["A", "B", "C"]
for i, j in zip(itertools.cycle(range(1, 4)), itertools.cycle(lst)):
    if i > 10:
        break
    else:
        print(i, j)

输出：

1 A
2 B
3 C
1 A
2 B
3 C
1 A
2 B
3 C

itertools.chain

itertools.chain(iterable1, iterable2, ..., iterableN) 函数将一个或多个迭代器合并成一个序列输出。

import itertools

lst1 = ["A", "B"]
lst2 = ["C", "D"]
for i in itertools.chain(lst1, lst2):
    print(i)

输出：

A
B
C
D

itertools.compress

itertools.compress(data, selectors) 函数根据 selectors 序列的 True 或 False 值对 data 序列进行过滤输出。

import itertools

lst = ["A", "B", "C", "D", "E"]
selectors = [False, True, True, False, True]
for i in itertools.compress(lst, selectors):
    print(i)

输出：

B
C
E

itertools.groupby

itertools.groupby(iterable, key=None) 函数类似于 SQL 语句中的 GROUP BY，将 iterable 中的元素根据 key 参数进行分组，返回一个包含 (key, group) 对构成的列表。key 是组的标识，group 是组成员构成的迭代器。

import itertools

lst = ["A", "B", "BC", "CC", "C"]
for k, g in itertools.groupby(lst, key=lambda x: x[0]):
    print(k, list(g))

输出：

A ['A']
B ['B', 'BC']
C ['CC', 'C']

itertools.permutations

itertools.permutations(iterable, r=None) 函数以 iterable 中元素的所有可能排列方式输出，r 参数用于控制排列元素的数量，如果不设置 r，则默认为 iterable 中所有元素的排列方式。

import itertools

lst = ["A", "B", "C"]
for i in itertools.permutations(lst, 2):
    print(i)

输出：

('A', 'B')
('A', 'C')
('B', 'A')
('B', 'C')
('C', 'A')
('C', 'B')

itertools.combinations

itertools.combinations(iterable, r) 函数以 iterable 中元素的所有可能的 r 元素组合方式输出。

import itertools

lst = ["A", "B", "C"]
for i in itertools.combinations(lst, 2):
    print(i)

输出：

('A', 'B')
('A', 'C')
('B', 'C')

总结

itertools 模块中的函数提供了处理迭代器的强大工具，这些函数既能用于提高程序的性能，也能让代码更加简洁易读。我们可以根据需要选用其中的某些函数，还可以根据需要组合使用。