公司动态
Python itertools 盘点真正用得上的20个常用方法
Python 标准库里藏着一个被严重低估的宝藏模块——itertools。几乎每个 Python 开发者都import过它但大多数人只用其中三四个函数剩下的就像工具箱里生锈的扳手从来没碰过。这篇报告把全部 20 个函数梳理清楚按照你在真实代码里会多频繁用到它来排序而不是照着官方文档的顺序念经。一、什么是itertools为什么值得认真学itertools是 Python 标准库中专门处理迭代器的模块灵感来自 APL、Haskell 和 SML 等函数式编程语言。它的核心哲学是构建一套迭代器代数——把简单的迭代器组合起来搭出复杂而高效的数据处理流水线。两个关键优势让它与众不同内存效率极高itertools的函数几乎全部是惰性求值lazy evaluation的也就是说数据不会一次性全部加载进内存而是用到哪块算哪块。处理一个几十 GB 的日志文件和处理一个 10 行的列表内存占用几乎没有区别。速度快底层用 C 实现比手写的 Python 循环快得多。研究表明itertools在节省内存的同时速度损耗微乎其微对绝大多数真实场景来说是非常划算的交换。理解itertools之前先要搞清楚两个概念可迭代对象Iterable能被for循环遍历的东西比如列表、字符串、字典。迭代器Iterator一种特殊的可迭代对象只能向前走不能回头调用next()一次吐出一个元素。itertools的所有函数返回的都是迭代器。二、函数全景图官方文档把 20 个函数分成三大类下面这张表是最清晰的全貌类别函数一句话描述通用迭代器chain,chain.from_iterable拼接多个可迭代对象通用迭代器islice切片迭代器通用迭代器compress按布尔掩码过滤通用迭代器filterfalse过滤掉满足条件的元素通用迭代器starmap解包参数后调用函数通用迭代器pairwise相邻两两配对3.10通用迭代器batched按批次分组3.12通用迭代器groupby按键分组需预排序通用迭代器accumulate累积计算如累加通用迭代器takewhile/dropwhile按条件截取/跳过通用迭代器tee复制迭代器通用迭代器zip_longest补齐长度的 zip无限迭代器count,cycle,repeat无限计数/循环/重复组合迭代器combinations,permutations,product,combinations_with_replacement排列组合笛卡尔积三、最常用的函数日常代码的主力军chain()和chain.from_iterable()— 拼接利器chain()是使用频率最高的函数之一把多个可迭代对象首尾相连变成一个统一的迭代器不需要真正合并成一个新列表。from itertools import chain # 把三个列表拼成一个流不占额外内存 result list(chain([1, 2], [3, 4], [5, 6])) # → [1, 2, 3, 4, 5, 6]chain.from_iterable()是它的变体专门处理列表的列表这种嵌套结构效果等同于chain(*nested)但不需要把整个嵌套列表先解包nested [[1, 2], [3, 4], [5, 6]] result list(chain.from_iterable(nested)) # → [1, 2, 3, 4, 5, 6]实际场景爬虫抓回来多页数据存在多个列表里用chain一行搞定合并比sum(lists, [])快得多。islice()— 迭代器的切片刀普通列表可以用[start:stop:step]切片但迭代器不行。islice()就是专门为迭代器设计的切片工具有一个重要限制不支持负数索引。from itertools import islice # 从无限计数器里取前5个 from itertools import count first_five list(islice(count(10), 5)) # → [10, 11, 12, 13, 14] # 跳过前3行读取第4到第10行 with open(bigfile.txt) as f: lines list(islice(f, 3, 10))处理大文件时这个函数特别好用——不用把整个文件读进内存按需取用。groupby()— 分组神器但有个大坑groupby()按照某个键把连续相同的元素分成组非常适合做数据聚合。但这里有个经典陷阱它只对连续相同的元素分组所以使用前必须先按同一个键排好序否则结果会出乎意料。from itertools import groupby data [ {name: Alice, dept: Engineering}, {name: Bob, dept: Engineering}, {name: Carol, dept: Marketing}, ] # 必须先排序 data.sort(keylambda x: x[dept]) for dept, members in groupby(data, keylambda x: x[dept]): print(dept, list(members))这个行为和 Unix 的uniq命令一模一样——只合并相邻的重复项。accumulate()— 不只是求和大多数人只知道accumulate能算累加但它接受任意二元函数可以做累积乘积、滚动最大值等各种操作from itertools import accumulate import operator nums [1, 2, 3, 4, 5] # 累加1 3 6 10 15 print(list(accumulate(nums))) # 累积乘积1 2 6 24 120 print(list(accumulate(nums, operator.mul))) # 滚动最大值1 2 3 4 5 print(list(accumulate(nums, max))) # 带初始值 print(list(accumulate(nums, initial100))) # → [100, 101, 103, 106, 110, 115]金融数据分析、游戏得分累计、滑动窗口统计这个函数都能派上用场。pairwise()和batched()— 两个为什么我以前要手写这个的函数Reddit 上那篇帖子特别点名了这两个函数因为太多人在不知道它们的情况下自己写了一堆笨拙的代码来实现同样的功能。pairwise()Python 3.10生成相邻元素的两两配对from itertools import pairwise steps list(pairwise([1, 2, 3, 4, 5])) # → [(1,2), (2,3), (3,4), (4,5)] # 计算相邻差值比如股价涨跌 prices [100, 105, 98, 110] changes [b - a for a, b in pairwise(prices)] # → [5, -7, 12]batched()Python 3.12把一个迭代器切成固定大小的批次最后一批可以比指定大小短from itertools import batched data range(10) for batch in batched(data, 3): print(batch) # → (0,1,2) (3,4,5) (6,7,8) (9,)批量调用 API、分页写数据库batched直接替代手写的[data[i:in] for i in range(0, len(data), n)]。四、进阶函数用得少但关键时刻救命starmap()— 优雅地解包参数当你有一个参数元组的列表想对每个元组调用同一个函数时starmap()比map()更优雅from itertools import starmap pairs [(2, 5), (3, 2), (10, 3)] result list(starmap(pow, pairs)) # → [32, 9, 1000] # 等价于但比下面这个更简洁 # [pow(a, b) for a, b in pairs]tee()— 复制迭代器但要小心tee(it, n)把一个迭代器复制成 n 个独立的迭代器每个都能独立遍历。听起来很美但有个内存陷阱它会在内部缓存所有已经被某个副本消费但其他副本还没读到的元素。如果两个副本的进度差很大内存消耗会很可观。from itertools import tee a, b tee([1, 2, 3, 4, 5]) # a 和 b 是独立的互不影响 print(list(a)) # → [1, 2, 3, 4, 5] print(list(b)) # → [1, 2, 3, 4, 5]原则tee之后原始迭代器就不要再用了。zip_longest()— 不丢数据的 zip内置zip()以最短的序列为准多余的元素直接丢掉。zip_longest()则以最长的为准短的用fillvalue补齐from itertools import zip_longest a [1, 2, 3, 4] b [x, y] result list(zip_longest(a, b, fillvalue-)) # → [(1,x), (2,y), (3,-), (4,-)]takewhile()和dropwhile()— 条件截断这两个函数是一对反义词都接受一个判断函数和一个序列from itertools import takewhile, dropwhile data [1, 4, 6, 3, 8] # 只要条件为真就继续取一旦为假立刻停止后面的不再检查 list(takewhile(lambda x: x 5, data)) # → [1, 4] # 跳过条件为真的从第一个为假的开始取所有剩余元素 list(dropwhile(lambda x: x 5, data)) # → [6, 3, 8]注意dropwhile的行为它跳过开头满足条件的元素但一旦开始输出后面的元素不再检查条件全部输出。五、无限迭代器小心别让程序跑死这三个函数会产生无穷无尽的元素必须配合islice()、takewhile()或zip()来截断否则程序会永远跑下去。from itertools import count, cycle, repeat # count从10开始步长2无限递增 # 10, 12, 14, 16, ... for n in count(10, 2): if n 20: break # cycle无限循环一个序列 # A, B, C, A, B, C, ... # 注意cycle 会把整个序列复制一份存在内存里 # repeat重复一个对象 n 次省略 n 则无限 list(repeat(7, 3)) # → [7, 7, 7]count和zip()配合是个经典技巧可以替代enumerate()而且步长更灵活。六、组合迭代器排列组合的数学世界这四个函数处理排列组合问题在密码学、测试用例生成、数学枚举场景里很常见。from itertools import combinations, permutations, product, combinations_with_replacement items [A, B, C] # 组合顺序无关不重复C(3,2)3种 list(combinations(items, 2)) # → [(A,B), (A,C), (B,C)] # 排列顺序有关不重复P(3,2)6种 list(permutations(items, 2)) # → [(A,B), (A,C), (B,A), (B,C), (C,A), (C,B)] # 笛卡尔积等价于嵌套 for 循环 list(product([0,1], repeat3)) # → 所有3位二进制数共8种 # 有放回组合允许重复元素 list(combinations_with_replacement(AB, 2)) # → [(A,A), (A,B), (B,B)]数学上这四者的关系可以用公式表示。对于从 n 个元素中取 r 个C(n,r)n!r!(n−r)!C(n,r)r!(n−r)!n!P(n,r)n!(n−r)!P(n,r)(n−r)!n!笛卡尔积的总数则是各序列长度的乘积。七、函数关系速查图用一张流程图展示这 20 个函数的分类关系八、实战组合使用才是真正的威力itertools的真正魔力在于组合。单个函数只是积木拼在一起才能搭出漂亮的东西。场景处理分页 API 数据跳过前两页取接下来三页的所有结果from itertools import chain, islice # 假设 fetch_page(n) 返回第 n 页的数据列表 pages (fetch_page(n) for n in range(10)) # 跳过前2页取接下来3页展平成一个流 result list(chain.from_iterable(islice(pages, 2, 5)))场景给数据流打上行号但从 1001 开始from itertools import count data [apple, banana, cherry] numbered list(zip(count(1001), data)) # → [(1001, apple), (1002, banana), (1003, cherry)]场景找出列表中所有连续递增的片段from itertools import groupby from itertools import pairwise nums [1, 2, 3, 1, 2, 4, 5] # 用 pairwise 计算差值再用 groupby 找连续的 1 段 diffs [b - a for a, b in pairwise(nums)] # → [1, 1, -2, 1, 2, 1]九、使用时的几个注意事项学到这里有几个容易踩的坑值得单独说清楚迭代器只能走一遍itertools返回的都是迭代器遍历一次就耗尽了。如果需要多次遍历要么转成list()要么用tee()。groupby必须预排序这是最常见的 bug 来源。groupby只合并相邻的相同元素不会自动排序。无限迭代器要有出口count()、cycle()没有终点一定要配合截断函数使用。cycle()会复制数据cycle内部会把整个序列存一份处理超大序列时要注意内存。islice不支持负索引想取最后 N 个元素用collections.deque(it, maxlenN)更合适。结语itertools就像一套精心设计的乐高积木——每一块单独看都简单但组合起来能搭出令人惊叹的结构。从最常用的chain、islice、groupby到很多人从未用过的pairwise、batched、starmap这 20 个函数覆盖了数据处理中绝大多数需要循环但又不想写一堆循环的场景。真正掌握itertools的标志不是能背出每个函数的参数而是在写代码时脑子里自然而然地想到这里可以用accumulate或者这里用chain.from_iterable比列表推导式更优雅。那种感觉值得花时间练习。