程序员求职经验分享与学习资料整理平台

网站首页 > 文章精选 正文

【Python神器】你可能忽视了的 itertools 模块

balukai 2025-07-28 15:17:12 文章精选 6 ℃

前言:为什么你需要掌握 itertools?

在日常 Python 编程中,我们常会面临 性能瓶颈、组合爆炸、懒加载问题、重复结构操作 等难题。很多人下意识去写循环或堆栈,却忽略了一个隐藏在 Python 标准库中的“性能之王”——itertools。

它是 Python 内置的函数式迭代处理库,不仅代码简洁优雅,还能极大提升执行效率。你也许用过 product 和 permutations,但远远没用到它的全部威力。

今天,我们就像阅读一篇专业文献那样,深入浅出地系统拆解 itertools 的全部精华,让你学完之后爱不释手。


模块全景图

我们将 itertools 按功能分成 4 类:

功能类别

函数名

无限迭代

count()、cycle()、repeat()

消耗迭代器

accumulate()、chain()、compress()、dropwhile()、takewhile()、filterfalse()

组合构造

product()、permutations()、combinations()、
combinations_with_replacement()

实用工具

tee()、zip_longest()、islice()、starmap()


一、无限迭代器:懒加载的代表作

1. itertools.count(start=0, step=1)

创建一个无限递增的数字序列。

from itertools import count
for i in count(10, 2):  # 输出 10, 12, 14...
    print(i)
    if i > 20:
        break

适用于生成唯一ID、时间戳序列、分页索引等。


2. itertools.cycle(iterable)

无限循环一个已有序列。

from itertools import cycle
colors = cycle(['red', 'green', 'blue'])
[next(colors) for _ in range(6)]
# 输出: ['red', 'green', 'blue', 'red', 'green', 'blue']

3. itertools.repeat(object, times=None)

重复某个对象(可控次数),是 map() 的最佳搭档。

from itertools import repeat
list(map(pow, range(5), repeat(2)))  # 相当于 [x**2 for x in range(5)]

二、组合生成器:排列组合不求人

4. product(*iterables, repeat=1)

笛卡尔积,适合遍历所有可能配置。

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

应用于网格搜索(Grid Search)、多参数测试。


5. permutations(iterable, r=None)

生成r个元素的全排列(不重复,顺序有关)。

from itertools import permutations
list(permutations('ABC', 2))  # [('A','B'), ('A','C'),...]

6. combinations(iterable, r)

从序列中取出 r 个不重复元素(顺序无关)。

from itertools import combinations
list(combinations('ABC', 2))  # [('A','B'), ('A','C'), ('B','C')]

7. combinations_with_replacement(iterable, r)

允许重复元素组合,比如 [1, 1]。

list(combinations_with_replacement('AB', 2))
# [('A', 'A'), ('A', 'B'), ('B', 'B')]

三、高级工具函数

8. accumulate(iterable, func=operator.add)

累积值生成器,默认做加法。

from itertools import accumulate
import operator
list(accumulate([1, 2, 3, 4], operator.mul))
# 输出: [1, 2, 6, 24](阶乘)

9. chain(*iterables)

将多个可迭代对象首尾相接。

from itertools import chain
list(chain([1, 2], [3, 4]))  # [1, 2, 3, 4]

10. compress(data, selectors)

类似 filter(),但使用一个布尔选择器列表。

from itertools import compress
data = 'ABCDEF'
selectors = [1, 0, 1, 0, 0, 1]
list(compress(data, selectors))  # ['A', 'C', 'F']

11. filterfalse(predicate, iterable)

只保留 predicate 为 False 的元素。

from itertools import filterfalse
list(filterfalse(lambda x: x % 2, range(10)))  # 偶数

12. dropwhile() 和 takewhile()

从一个序列中“跳过”或“取出”前面满足某条件的部分。

from itertools import dropwhile, takewhile
nums = [1, 4, 6, 2, 1]
list(dropwhile(lambda x: x < 5, nums))  # [6, 2, 1]
list(takewhile(lambda x: x < 5, nums))  # [1, 4]

四、冷门但神奇的辅助工具

13. tee(iterable, n=2)

复制一个迭代器为多个“独立快照”。

from itertools import tee
a, b = tee(range(3))
list(a)  # [0,1,2]
list(b)  # [0,1,2]

注意:它会缓存数据,内存敏感时慎用。


14. islice(iterable, start, stop, step)

切片操作,但作用于迭代器。

from itertools import islice
list(islice(range(10), 2, 8, 2))  # [2, 4, 6]

15. starmap(func, iterable)

用于函数参数“拆包”的 map。

from itertools import starmap
list(starmap(pow, [(2, 5), (3, 2), (10, 3)]))  # [32, 9, 1000]

真实案例:构造懒加载组合推荐系统

假设你正在设计一个电影推荐系统,为了生成 每对电影之间的组合打分

from itertools import combinations
movies = ['Inception', 'Interstellar', 'Memento', 'Tenet']
pairs = combinations(movies, 2)

for a, b in pairs:
    print(f"{a} vs {b}: score = {hash(a) % 100 + hash(b) % 100}")

你可以再配合 starmap、filterfalse 做出更复杂的推荐逻辑,而不需要全量装入内存!


总结 & 实战建议

技术要点

场景

product, permutations

模型网格搜索、数据增强

chain, accumulate

日志合并、滚动计算

islice, tee

分页、分批训练

cycle, repeat

数据增强、Mock 数据生成

compress, filterfalse

权重选择、特征筛选


延伸阅读

  • 官方文档:itertools — functions creating iterators for efficient looping
  • Python Cookbook, 3rd Edition(David Beazley & Brian K. Jones)

结语

itertools 不仅是工具,更是一种写 Python 的思维方式。用生成器思维构建高效、可复用的代码结构,是高级工程师的标志之一。

下次你想写 for-loop 套 for-loop 前,先想想:有没有 itertools 能帮你一次搞定?

最近发表
标签列表