ZetCode

Python itertools 模块

最后修改于 2025 年 4 月 2 日

itertools 模块提供了一套用于处理迭代器的高速、内存高效的工具。这些函数受到函数式编程语言中构造的启发,并且旨在与 Python 的迭代器协议无缝协作。本指南涵盖了所有 itertools 函数,并附有实际示例、性能考量和真实应用。

无限迭代器

itertools 提供了三个用于创建无限迭代器的函数:countcyclerepeat。这些函数无限生成值,直到显式停止。本示例演示了它们的基本用法模式和常见应用。

infinite_iterators.py
import itertools

# 1. count(start=0, step=1) - infinite arithmetic sequence
counter = itertools.count(start=5, step=3)
print("Count:", [next(counter) for _ in range(5)])  # [5, 8, 11, 14, 17]

# 2. cycle(iterable) - infinitely cycle through an iterable
cycler = itertools.cycle('ABC')
print("Cycle:", [next(cycler) for _ in range(6)])  # ['A', 'B', 'C', 'A', 'B', 'C']

# 3. repeat(object[, times]) - repeat object indefinitely or fixed times
repeater = itertools.repeat('hello', 3)
print("Repeat:", list(repeater))  # ['hello', 'hello', 'hello']

# Additional example: Creating sliding windows with count
data = [10, 20, 30, 40, 50]
windows = zip(itertools.count(), data, data[1:], data[2:])
print("Sliding windows:", list(windows))
# [(0, 10, 20, 30), (1, 20, 30, 40), (2, 30, 40, 50)]

count 生成带有可选开始值和步长的无限数字序列。cycle 无限重复有限可迭代对象的元素。repeat 产生相同的值,可以是无限的,也可以是指定次数。

这些无限迭代器是内存高效的,因为它们按需生成值。它们经常与 zipislice 一起使用来创建有限序列,或者与需要无限值流的函数一起使用。

组合迭代器

组合迭代器(productpermutationscombinations 等)从输入可迭代对象生成复杂的序列。这些函数在解决涉及组合、排列或笛卡尔积的问题时非常有用。

combinatoric_iterators.py
import itertools

# 4. product(*iterables, repeat=1) - Cartesian product
dice = itertools.product([1, 2, 3], ['a', 'b'])
print("Product:", list(dice))
# [(1, 'a'), (1, 'b'), (2, 'a'), (2, 'b'), (3, 'a'), (3, 'b')]

# 5. permutations(iterable, r=None) - r-length permutations
letters = itertools.permutations('ABC', 2)
print("Permutations:", list(letters))
# [('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'C'), ('C', 'A'), ('C', 'B')]

# 6. combinations(iterable, r) - r-length combinations, no repeats
cards = itertools.combinations(['♥A', '♦K', '♣Q'], 2)
print("Combinations:", list(cards))
# [('♥A', '♦K'), ('♥A', '♣Q'), ('♦K', '♣Q')]

# 7. combinations_with_replacement(iterable, r) - with repeats
dice_rolls = itertools.combinations_with_replacement([1, 2, 3], 2)
print("Combinations w/replacement:", list(dice_rolls))
# [(1, 1), (1, 2), (1, 3), (2, 2), (2, 3), (3, 3)]

# Additional example: Generating truth tables
variables = [False, True]
truth_table = itertools.product(variables, repeat=2)
print("Truth table:")
for a, b in truth_table:
    print(f"{a} AND {b} = {a and b}")

product 计算输入可迭代对象的笛卡尔积,等同于嵌套的 for 循环。permutations 生成所有可能的排序,不重复元素。combinations 生成顺序无关紧要的子序列,而 combinations_with_replacement 允许重复元素。

这些函数在概率、统计、游戏开发和算法设计中特别有用。它们可以生成大量的中间结果,因此通常与其他 itertools 函数一起使用来限制输出。

在最短输入时终止的迭代器

此组包括 chainzip_longestfilterfalse 等函数,它们处理多个可迭代对象,直到最短的可迭代对象耗尽(zip_longest 除外)。这些对于处理多个数据流至关重要。

terminating_iterators.py
import itertools

# 8. chain(*iterables) - concatenate iterables
merged = itertools.chain('ABC', [1, 2, 3], (True, False))
print("Chain:", list(merged))
# ['A', 'B', 'C', 1, 2, 3, True, False]

# 9. zip_longest(*iterables, fillvalue=None) - zip to longest iterable
names = ['Alice', 'Bob']
scores = [85, 92, 78]
zipped = itertools.zip_longest(names, scores, fillvalue='N/A')
print("Zip longest:", list(zipped))
# [('Alice', 85), ('Bob', 92), ('N/A', 78)]

# 10. filterfalse(predicate, iterable) - elements where predicate is False
numbers = [0, 1, 0, 2, 3, 0, 4]
non_zeros = itertools.filterfalse(lambda x: x == 0, numbers)
print("Filterfalse:", list(non_zeros))  # [1, 2, 3, 4]

# 11. islice(iterable, stop) or islice(iterable, start, stop[, step]) - slice iterator
infinite = itertools.count()
first_5_evens = itertools.islice(infinite, 0, 10, 2)
print("Islice:", list(first_5_evens))  # [0, 2, 4, 6, 8]

# Additional example: Processing batches
data = range(100)
batch_size = 10
for batch in itertools.islice(data, 0, None, batch_size):
    print("Batch:", list(itertools.islice(data, batch, batch + batch_size)))

chain 对于组合不同的数据源特别有用。zip_longest 可以优雅地处理不等长的可迭代对象。filterfalse 提供了内置 filter 的反向操作。islice 能够高效地对迭代器进行切片,而无需将其转换为列表。

这些函数在数据处理管道中大放异彩,您可以在其中组合、过滤或分块数据流,而无需将所有数据加载到内存中。它们经常与文件处理和数据库查询一起使用。

分组和过滤

groupbytakewhile/dropwhile 函数提供了强大的工具来组织和过滤顺序数据。这些对于数据分析和预处理任务尤其有价值。

grouping_filtering.py
import itertools

# 12. groupby(iterable, key=None) - group consecutive elements
animals = ['ant', 'bee', 'cat', 'dog', 'eagle', 'flamingo']
grouped = itertools.groupby(animals, key=lambda x: x[0])
print("Groupby:")
for key, group in grouped:
    print(f"{key}: {list(group)}")
# a: ['ant']
# b: ['bee']
# c: ['cat']
# d: ['dog']
# e: ['eagle']
# f: ['flamingo']

# 13. takewhile(predicate, iterable) - take until predicate fails
numbers = [1, 4, 6, 8, 2, 5, 3]
taken = itertools.takewhile(lambda x: x < 7, numbers)
print("Takewhile:", list(taken))  # [1, 4, 6]

# 14. dropwhile(predicate, iterable) - drop until predicate fails
dropped = itertools.dropwhile(lambda x: x < 7, numbers)
print("Dropwhile:", list(dropped))  # [8, 2, 5, 3]

# Additional example: Processing log files
log_lines = [
    "INFO: System started",
    "INFO: User logged in",
    "ERROR: File not found",
    "INFO: Request processed",
    "ERROR: Database timeout"
]

# Group by log level
get_level = lambda line: line.split(':')[0]
for level, lines in itertools.groupby(log_lines, key=get_level):
    print(f"\n{level} messages:")
    for line in lines:
        print("  ", line.split(':', 1)[1].strip())

groupby 按键对连续的相同元素进行分组(需要已排序的输入才能完全分组)。takewhile 生成项,直到谓词失败,而 dropwhile 跳过项,直到谓词失败,然后生成其余项。

这些函数对于处理顺序数据(如日志、时间序列或任何分组记录)非常有价值。它们能够进行高效处理,而无需将整个数据集加载到内存中。

性能注意事项

虽然 itertools 函数是内存高效的,但它们的性能特征各不相同。本节比较了常见操作,并演示了处理大型数据集的优化技术。

performance.py
import itertools
import timeit
import random

# 15. Comparing chain methods
def test_chain():
    list(itertools.chain(range(1000), range(1000, 2000)))

def test_concat():
    list(range(1000)) + list(range(1000, 2000))

print("Chain vs concat:")
print("itertools.chain:", timeit.timeit(test_chain, number=10000))
print("list concatenation:", timeit.timeit(test_concat, number=10000))

# 16. Memory efficiency demonstration
large_range = itertools.count()  # Infinite, uses almost no memory
# Compare with list(range(1000000)) which would consume significant memory

# 17. Early termination with islice
def process_data():
    data = itertools.count()  # Infinite stream
    processed = (x**2 for x in itertools.islice(data, 1000000))
    return sum(processed)  # Doesn't store all squared values

print("\nProcessing 1M numbers:", process_data())

# Additional example: Filtering large datasets
def large_dataset():
    return (random.random() for _ in range(1000000))

# Memory-efficient filtering
positive = itertools.filterfalse(lambda x: x < 0.5, large_dataset())
print("\nCount > 0.5:", sum(1 for _ in itertools.islice(positive, 0, 100000)))

基准测试显示,对于大型可迭代对象,itertools.chain 比列表连接更快。内存效率示例演示了 itertools 如何处理理论上无限的序列。提前终止示例在不将其物化在内存中的情况下处理大型范围。

关键要点:itertools 函数在内存效率和惰性求值方面表现出色。当处理大型或无限序列时,它们特别有利,但对于小型数据集,内置函数可能更简单且性能同样出色。

实际应用

这些示例演示了 itertools 在常见编程场景中的实际应用,从数据分析到算法实现。

applications.py
import itertools
import operator

# 18. Running averages
def running_avg(data):
    it = itertools.accumulate(data, operator.add)
    for i, total in enumerate(it, 1):
        yield total / i

print("Running averages:", list(running_avg([10, 20, 30, 40])))
# [10.0, 15.0, 20.0, 25.0]

# 19. Pairwise iteration (Python 3.10+ has itertools.pairwise)
def pairwise(iterable):
    a, b = itertools.tee(iterable)
    next(b, None)
    return zip(a, b)

print("Pairwise differences:", [(y-x) for x, y in pairwise([1, 3, 6, 10])])
# [2, 3, 4]

# 20. Pagination with islice
def paginate(items, page_size):
    page_start = 0
    while True:
        page = list(itertools.islice(items, page_start, page_start + page_size))
        if not page:
            break
        yield page
        page_start += page_size

data = range(0, 10)
print("Paginated data:")
for page in paginate(data, 3):
    print(page)
# [0, 1, 2]
# [3, 4, 5]
# [6, 7, 8]
# [9]

# Additional example: Cartesian product for parameter grids
params = {
    'learning_rate': [0.01, 0.1],
    'batch_size': [32, 64],
    'optimizer': ['adam', 'sgd']
}

param_grid = itertools.product(*params.values())
print("\nParameter combinations:")
for combo in param_grid:
    print(dict(zip(params.keys(), combo)))

运行平均值示例显示了 accumulate 如何简化有状态的计算。成对迭代演示了时间序列分析中的常见模式。分页示例说明了如何分块处理大型数据集。参数网格示例在机器学习超参数调优中非常有用。

这些模式广泛应用于数据处理、科学计算和 Web 开发。itertools 函数有助于保持代码简洁和内存高效。

最佳实践

使用 itertools 来高效处理大型或无限序列。组合多个 itertools 函数以创建复杂的处理管道。对于常见的迭代模式,优先使用 itertools 而不是手动实现。请记住,许多 itertools 函数会消耗迭代器(例如 tee),因此它们无法重复使用。为方便维护,请记录复杂的 itertools 管道。在可读性更好的简单情况下,请考虑生成器表达式。

资料来源

从这些资源中了解更多:Python itertools 文档more-itertools 库

作者

我的名字是 Jan Bodnar,我是一位充满热情的程序员,拥有丰富的编程经验。我自 2007 年以来一直在撰写编程文章。迄今为止,我已撰写了 1400 多篇文章和 8 本电子书。我在教授编程方面拥有十多年的经验。

列出所有 Python 教程