Python 基础体系 · 第 43/112 篇。示例统一以 Python 3.14 为语言基线;第三方库使用与其兼容的现代稳定版本,版本敏感行为会单独说明。

Python itertools:惰性组合、无限序列、分组和内存边界

itertools 不是一组“更短的循环写法”,而是一套围绕迭代器建立的组合式工具。它把数据处理拆成若干可以串联的步骤:生成、变换、过滤、截取、分组、批处理和组合。

这些工具共同遵循一个重要方向:

尽量延迟计算,只在下游请求下一个元素时推进上游,并避免一次性构造完整结果。

但“惰性”不等于“零内存”,“返回迭代器”也不等于“可以处理无限输入”。cycle()tee()product()combinations()groupby() 都有不同的缓存或消费边界。理解这些边界,才能判断一个流水线是否真的适合生产环境。


1. 先建立模型:可迭代对象、迭代器和惰性

1.1 可迭代对象不一定是迭代器

可迭代对象是可以交给 iter() 获取迭代器的对象:

values = [10, 20, 30]

iterator = iter(values)

print(next(iterator))  # 10
print(next(iterator))  # 20
print(next(iterator))  # 30

迭代器至少提供两个行为:

  1. __iter__() 返回自身;
  2. __next__() 返回下一个元素,耗尽后抛出 StopIteration

可以把它形式化为:

Inext(x,I)I \xrightarrow{\text{next}} (x, I')

其中:

  • II 是当前迭代器状态;
  • xx 是本次产生的元素;
  • II' 是推进后的状态。

当没有更多元素时:

InextStopIterationI \xrightarrow{\text{next}} \text{StopIteration}

列表通常是可迭代对象,但不是迭代器;iter(values) 返回的对象才是迭代器。

values = [1, 2, 3]

print(iter(values) is values)  # False

iterator = iter(values)
print(iter(iterator) is iterator)  # True

这个区别决定了一个对象是否具有“游标状态”。列表可以反复创建新的迭代器,而迭代器本身会被消耗。

values = [1, 2, 3]

print(list(values))  # [1, 2, 3]
print(list(values))  # [1, 2, 3]

iterator = iter(values)

print(list(iterator))  # [1, 2, 3]
print(list(iterator))  # []

itertools 中的大多数函数接收可迭代对象,但返回的是有状态的迭代器。Python 官方文档将这些工具描述为用于构造高效循环的迭代器构建块,并强调它们可以组合成“迭代器代数”。(docs.python.org)


1.2 惰性意味着“按需推进”

考虑下面的代码:

from itertools import islice

def numbers():
    for number in range(10):
        print(f"生成 {number}")
        yield number

result = islice(numbers(), 2, 5)

print("迭代器已创建")
print(next(result))
print(next(result))

输出类似:

迭代器已创建
生成 0
生成 1
生成 2
2
生成 3
3

创建 islice 时,numbers() 并没有立即执行。第一次调用 next(result) 时,上游被推进到下标 2;第二次调用时才继续取得 3

可以把迭代器流水线写成:

S0T1T2CS_0 \rightarrow T_1 \rightarrow T_2 \rightarrow \cdots \rightarrow C

其中:

  • S0S_0 是源数据;
  • TiT_i 是一个惰性转换;
  • CC 是最终消费者,例如 fornext()list() 或文件写入。

只有消费者请求元素时,请求才从右向左传播,最终促使源数据产生足够的输入。

因此:

pipeline = map(str.upper, filter(str.isalpha, data))

通常不会立即处理 data。而:

result = list(pipeline)

会驱动整个流水线,直到输入耗尽。


1.3 惰性不等于完全不占内存

一个迭代器可能只保留少量状态,也可能为了未来输出保存大量数据。

例如:

  • map() 通常只保存函数和上游迭代器;
  • islice() 通常只保存切片位置;
  • cycle() 会保存已经读过的元素;
  • tee() 会缓存较快分支已经消费、较慢分支尚未消费的元素;
  • product()permutations()combinations() 会先把输入转换为池;
  • groupby() 返回的每个分组迭代器与主迭代器共享底层输入。

所以判断内存复杂度时,不能只看返回类型是否为 iterator,而要问:

这个工具为了产生未来的结果,必须记住哪些已经消费的数据?


2. itertools 的基本分类

Python 3.14 的 itertools 主要可以按数据流性质分为四类:

  1. 逐元素变换accumulate()compress()filterfalse()starmap()
  2. 结构变换与截取chain()batched()islice()pairwise()takewhile()dropwhile()zip_longest()
  3. 状态或无限序列count()cycle()repeat()
  4. 组合生成product()permutations()combinations()combinations_with_replacement()
  5. 分组和分支groupby()tee()

所有函数都不应该被机械地视为“无内存开销工具”。官方文档明确指出,一些工具会产生无限长度流,必须由循环或其他工具截断;同时,product() 会在开始生成结果前完全消费输入,cycle() 可能需要大量辅助存储。(docs.python.org)


3. 无限序列:count()repeat()cycle()

3.1 count():等差序列生成器

count(start=0, step=1) 产生:

an=start+n×stepa_n = start + n \times step

代码如下:

from itertools import count, islice

values = islice(count(10, 3), 5)

print(list(values))

输出:

[10, 13, 16, 19, 22]

count() 本身没有终点:

from itertools import count

for number in count(1):
    print(number)

这个循环不会自然结束,必须通过 break 或组合一个有限截断器:

from itertools import count, islice

first_five = islice(count(1), 5)

for number in first_five:
    print(number)

输出:

1
2
3
4
5

count() 的数据流是:

start
  ↓ + step
start + step
  ↓ + step
start + 2 * step
  ↓
...

它只需要保存当前值,因此自身的辅助空间是常量级。但无限长度意味着终止责任转移给调用者


3.2 浮点数的边界

浮点数累加可能产生舍入误差:

from itertools import count, islice

print(list(islice(count(0.0, 0.1), 5)))

可能得到:

[0.0, 0.1, 0.2, 0.30000000000000004, 0.4]

官方文档建议,在需要更好精度时,使用乘法形式:

from itertools import count, islice

values = (
    0.0 + 0.1 * index
    for index in islice(count(), 5)
)

print(list(values))

不过,乘法并不意味着浮点数就获得了十进制定点语义。如果业务要求金额精确到分,应使用整数分或 decimal.Decimal,而不是依赖 count() 的浮点输出。


3.3 repeat():重复同一个对象

repeat(value, times=None) 反复产生同一个对象:

from itertools import repeat

print(list(repeat("ready", 3)))

输出:

['ready', 'ready', 'ready']

它经常与 map() 配合,为每个输入提供同一个参数:

from itertools import repeat

powers = map(pow, range(5), repeat(2))

print(list(powers))

输出:

[0, 1, 4, 9, 16]

这里的逻辑等价于:

result = []

for value in range(5):
    result.append(pow(value, 2))

repeat(2) 不会构造 [2, 2, 2, 2, ...]。它只是按需返回同一个对象,因此适合为有限输入补充常量参数。

要注意“重复对象”和“复制对象”的区别:

from itertools import repeat

items = list(repeat([], 2))
items[0].append("x")

print(items)

输出:

[['x'], ['x']]

两个位置引用的是同一个列表对象。repeat() 重复的是对象引用,不会为每次输出执行深拷贝。


3.4 cycle():循环输入,但需要缓存输入

cycle(iterable) 会先遍历输入,并保存每个元素;输入耗尽后,再从保存的数据中无限循环。官方文档的近似实现正是“边输出边追加到 saved,输入结束后反复遍历 saved”。(docs.python.org)

from itertools import cycle, islice

colors = islice(cycle(["red", "green", "blue"]), 8)

print(list(colors))

输出:

['red', 'green', 'blue', 'red', 'green', 'blue', 'red', 'green']

如果输入长度为 nncycle() 至少需要保存这 nn 个元素:

辅助空间=O(n)\text{辅助空间} = O(n)

这和 count() 不同:

from itertools import count, islice

# count() 不需要保存已经产生的所有整数
print(list(islice(count(), 1_000_000)))

而:

from itertools import cycle

# cycle() 必须记住输入的全部元素,才能再次循环
result = cycle(large_input)

如果输入本身是无限迭代器,cycle() 会永远停留在“读取并缓存第一轮输入”的阶段,不会进入第二轮:

from itertools import count, cycle

# 不会产生任何可观察的循环输出;
# cycle() 会不断缓存 count() 的新值。
result = cycle(count())

因此:

  • 对有限、小型输入,可以使用 cycle()
  • 对无限输入不能使用 cycle()
  • 对大型输入,优先考虑显式状态机或有界循环,而不是让 cycle() 无限增长缓存。

4. 截取与边界:islice()takewhile()dropwhile()

4.1 islice() 是迭代器版本的切片

islice(iterable, start, stop, step) 类似序列切片,但不支持负数索引、负数停止位置或非正步长。它不会把整个输入转换成列表。(docs.python.org)

from itertools import islice

data = "ABCDEFG"

print(list(islice(data, 2)))          # ['A', 'B']
print(list(islice(data, 2, 5)))       # ['C', 'D', 'E']
print(list(islice(data, 0, None, 2))) # ['A', 'C', 'E', 'G']

它适用于无限流:

from itertools import count, islice

first_ten = islice(count(100), 10)

print(list(first_ten))

输出:

[100, 101, 102, 103, 104, 105, 106, 107, 108, 109]

但是,islice()step 只决定输出哪些元素,不代表上游只推进输出数量。假设:

from itertools import count, islice

source = count()
selected = islice(source, 0, 10, 3)

print(list(selected))  # [0, 3, 6, 9]
print(next(source))    # 10

为了得到 0、3、6、9,上游实际上已经推进到 10。官方文档说明,完全消费 islice 后,如果输入是迭代器,上游会被推进 max(start, stop) 个位置,而不是仅推进实际输出的元素个数。(docs.python.org)


4.2 takewhile() 的分界元素会被消费

takewhile(predicate, iterable) 持续输出满足条件的元素,一旦遇到第一个不满足条件的元素就停止:

from itertools import takewhile

data = [2, 4, 6, 7, 8]

print(list(takewhile(lambda x: x % 2 == 0, data)))

输出:

[2, 4, 6]

关键边界是:7 已经被读取并用于判断,因此它从输入迭代器中消失了。

from itertools import takewhile

source = iter([2, 4, 6, 7, 8])
prefix = takewhile(lambda x: x % 2 == 0, source)

print(list(prefix))  # [2, 4, 6]
print(next(source))  # 8

如果需要同时保留“满足条件的前缀”和“第一个失败元素及其之后的剩余输入”,不能直接用 takewhile()。可以显式写一个带回退逻辑的生成器:

def split_at_first_failure(iterable, predicate):
    iterator = iter(iterable)

    for value in iterator:
        if not predicate(value):
            return value, iterator
        yield value

source = iter([2, 4, 6, 7, 8])
prefix = split_at_first_failure(source, lambda x: x % 2 == 0)

print(list(prefix))  # [2, 4, 6]
failed = next(prefix, None)

不过这个简单版本无法通过普通 yield 同时自然返回失败元素;实际工程中通常会把“前缀”和“剩余流”封装成明确的数据结构,或使用 more-itertools 提供的专用工具。核心事实不变:takewhile() 的停止条件依赖于读取那个失败元素。(docs.python.org)


4.3 dropwhile() 可能有较长启动延迟

dropwhile(predicate, iterable) 丢弃开头满足条件的元素,一旦谓词第一次返回假值,就输出当前元素以及之后所有元素:

from itertools import dropwhile

data = [1, 3, 5, 8, 2, 4]

print(list(dropwhile(lambda x: x < 5, data)))

输出:

[5, 8, 2, 4]

它与普通过滤器不同。dropwhile() 只对开头连续区间应用谓词;一旦找到第一个不满足条件的元素,之后不再过滤。

from itertools import dropwhile

data = [1, 3, 5, 2, 7, 4]

print(list(dropwhile(lambda x: x < 5, data)))

输出:

[5, 2, 7, 4]

如果需求是“过滤掉所有小于 5 的元素”,应使用 filterfalse()

from itertools import filterfalse

print(list(filterfalse(lambda x: x < 5, data)))

输出:

[5, 7]

dropwhile() 在前缀很长时不会产生任何输出,直到找到第一个失败元素;这也是它可能产生启动延迟的原因。(docs.python.org)


5. 批处理:batched() 与输入消费边界

5.1 batched() 的基本语义

batched(iterable, n) 将输入分成长度最多为 n 的元组:

from itertools import batched

data = "ABCDEFG"

print(list(batched(data, 3)))

输出:

[('A', 'B', 'C'), ('D', 'E', 'F'), ('G',)]

它只为当前批次保存最多 n 个元素:

辅助空间=O(n)\text{辅助空间} = O(n)

而不是一次性保存整个输入。

batched() 在 Python 3.12 中加入,Python 3.13 增加了 strict 参数。(docs.python.org)


5.2 strict=True 将数据完整性变成显式约束

如果每个批次都必须完整,可以使用:

from itertools import batched

data = range(7)

try:
    for batch in batched(data, 3, strict=True):
        print(batch)
except ValueError as exc:
    print(type(exc).__name__, exc)

输出:

(0, 1, 2)
(3, 4, 5)
ValueError batched(): incomplete batch

注意,最后一个不完整批次是在被取出时发现的。前面的完整批次已经交给调用者处理,因此如果批处理对应事务、文件分片或网络请求,异常处理必须考虑:

  1. 前面的批次是否已经产生副作用;
  2. 最后一个批次失败后是否需要回滚;
  3. 是否应先写入临时位置,再进行提交;
  4. 重试时是否会重复处理前面的批次。

strict=True 只能检测批次边界是否完整,不能保证元素内容有效,也不能保证处理过程具备事务性。


5.3 逐步观察批处理的惰性

from itertools import batched

def source():
    for value in range(5):
        print(f"读取 {value}")
        yield value

batches = batched(source(), 2)

print("创建完成")
print(next(batches))
print(next(batches))

输出:

创建完成
读取 0
读取 1
(0, 1)
读取 2
读取 3
(2, 3)

创建 batches 不会读取输入。第一次 next() 恰好读取两个元素并立即产生一个批次;不会为了知道后面还有多少数据而扫描完整输入。


6. 组合工具:结果惰性,但输入和结果空间都可能巨大

组合工具最容易产生误判:

组合结果是惰性产生的,不代表组合问题本身规模很小。

设输入有 nn 个元素。

6.1 product():笛卡尔积

product(A, B) 等价于嵌套循环:

from itertools import product

print(list(product([1, 2], ["a", "b"])))

输出:

[(1, 'a'), (1, 'b'), (2, 'a'), (2, 'b')]

如果各输入长度为 n1,n2,,nkn_1,n_2,\ldots,n_k,结果数量为:

N=i=1kniN = \prod_{i=1}^{k} n_i

当使用 product(values, repeat=r),且 values 长度为 nn 时:

N=nrN = n^r

例如:

from itertools import product

for item in product("AB", repeat=3):
    print("".join(item))

输出顺序为:

AAA
AAB
ABA
ABB
BAA
BAB
BBA
BBB

product() 的输出是惰性的,但开始运行前会完全消费每个输入,并保留输入池。官方文档因此指出它只适用于有限输入。(docs.python.org)

这意味着下面的代码无法工作:

from itertools import count, product

# product() 需要先完全消费 count(),因此不会产生第一个结果
result = product(count(), ["x", "y"])

product() 的内存边界与结果数量不同:

  • 输入池内存:与输入元素数量有关;
  • 输出结果内存:如果逐个消费,可以保持较低;
  • 总计算时间:仍可能达到乘法级别。

因此可以安全地做:

from itertools import product, islice

candidates = product(range(1000), repeat=3)

for candidate in islice(candidates, 10):
    print(candidate)

这里只取前 10 个结果,但 product() 仍然需要先把 range(1000) 转换为池;它没有生成全部 100031000^3 个元组。


6.2 permutations():有顺序的无重复选择

permutations(iterable, r)nn 个位置不同的元素中选择长度为 rr 的排列,数量为:

P(n,r)=n!(nr)!P(n,r)=\frac{n!}{(n-r)!}

from itertools import permutations

print(list(permutations("ABC", 2)))

输出:

[('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'C'), ('C', 'A'), ('C', 'B')]

permutations("ABC", 2)combinations("ABC", 2) 的差别在于顺序:

permutations: AB AC BA BC CA CB
combinations: AB AC BC

输入元素按照位置区分,而不是按照值去重:

from itertools import permutations

print(list(permutations(["A", "A"], 2)))

输出:

[('A', 'A'), ('A', 'A')]

两个 A 来自不同位置,所以会生成两个排列,即使打印出来的值相同。


6.3 combinations():无顺序且不重复选择

combinations(iterable, r) 选择 rr 个不同位置的元素,不关心选择顺序。结果数量为:

C(n,r)=n!r!(nr)!C(n,r)=\frac{n!}{r!(n-r)!}

from itertools import combinations

print(list(combinations("ABCD", 2)))

输出:

[('A', 'B'), ('A', 'C'), ('A', 'D'), ('B', 'C'), ('B', 'D'), ('C', 'D')]

其核心状态可以理解为递增下标:

初始: [0, 1]       -> AB
推进: [0, 2]       -> AC
推进: [0, 3]       -> AD
推进: [1, 2]       -> BC
推进: [1, 3]       -> BD
推进: [2, 3]       -> CD
结束

官方文档给出的近似算法正是通过下标数组寻找从右侧开始还能增加的位置,并将其后的下标重新设置为连续递增值。(docs.python.org)

combinations() 会先把输入转换为元组池,因此输入本身必须是有限的,并且需要额外保存输入元素。


6.4 combinations_with_replacement():允许重复选择

这个工具允许同一个位置对应的元素被多次选中:

from itertools import combinations_with_replacement

print(list(combinations_with_replacement("ABC", 2)))

输出:

[('A', 'A'), ('A', 'B'), ('A', 'C'), ('B', 'B'), ('B', 'C'), ('C', 'C')]

数量为:

C(n+r1,r)=(n+r1)!r!(n1)!C(n+r-1,r)=\frac{(n+r-1)!}{r!(n-1)!}

例如从 3 种配料中选择 4 份,允许重复:

C(3+41,4)=C(6,4)=15C(3+4-1,4)=C(6,4)=15

它与 product() 的差别不是“是否允许值重复”这么简单,而是结果是否保持非递减位置顺序:

product("ABC", repeat=2):
AA AB AC BA BB BC CA CB CC

combinations_with_replacement("ABC", 2):
AA AB AC BB BC CC

6.5 组合工具的工程边界

即使使用:

for item in combinations(data, 2):
    process(item)

也不能消除组合数量。若 len(data)=100_000,二元组合数量为:

C(100000,2)=4,999,950,000C(100000,2)=4,999,950,000

这类任务的问题不是“列表是否一次性创建”,而是候选空间本身过大。需要在组合生成前加入约束,例如:

  • 先过滤输入;
  • 限制 r
  • 使用业务规则剪枝;
  • 只取前若干结果;
  • 将问题改写为索引、哈希连接或专用算法。

itertools 负责高效枚举,不负责把指数级或组合级问题变成线性问题。


7. 连接、补齐和参数展开:chain()zip_longest()starmap()

7.1 chain():多个输入串成一个流

from itertools import chain

parts = chain([1, 2], (3, 4), "56")

print(list(parts))

输出:

[1, 2, 3, 4, '5', '6']

chain() 依次消费每个输入,前一个耗尽后才访问下一个:

from itertools import chain

def source(name, values):
    for value in values:
        print(f"{name}: {value}")
        yield value

result = chain(
    source("A", [1, 2]),
    source("B", [3, 4]),
)

print(next(result))
print(next(result))
print(next(result))

输出顺序体现了数据流:

A: 1
1
A: 2
2
B: 3
3

如果输入本身是一个“可迭代对象的可迭代对象”,应使用 chain.from_iterable()

from itertools import chain

chunks = [[1, 2], [3], [4, 5]]

print(list(chain.from_iterable(chunks)))

输出:

[1, 2, 3, 4, 5]

chain.from_iterable() 对外层输入也是惰性读取,适合串联动态产生的多个数据源。


7.2 zip_longest():保留较长输入的尾部

内置 zip() 在最短输入耗尽时停止:

print(list(zip([1, 2, 3], ["a", "b"])))

输出:

[(1, 'a'), (2, 'b')]

zip_longest() 则使用填充值继续产生结果:

from itertools import zip_longest

result = zip_longest(
    [1, 2, 3],
    ["a", "b"],
    fillvalue=None,
)

print(list(result))

输出:

[(1, 'a'), (2, 'b'), (3, None)]

它的输出数量取决于最长输入。对无限输入要特别小心:

from itertools import count, zip_longest

# count() 无限,zip_longest() 也不会自然结束
result = zip_longest([1, 2], count())

如果必须处理无限输入,应明确使用 islice() 或其他终止条件。


7.3 starmap():输入已经是参数元组时展开调用

map(function, iterable) 相当于:

function(value)

starmap(function, iterable) 相当于:

function(*args)
from itertools import starmap

calls = [(2, 5), (3, 2), (10, 3)]

print(list(starmap(pow, calls)))

输出:

[32, 9, 1000]

如果使用 map(),需要额外写一个解包函数:

result = map(lambda args: pow(*args), calls)

starmap() 更直接表达了数据形状:输入的每个元素本身就是一次调用的参数集合。


8. 过滤与选择:compress()filterfalse() 和谓词边界

8.1 compress():用选择器筛选数据

from itertools import compress

data = ["A", "B", "C", "D"]
selectors = [True, False, True, False]

print(list(compress(data, selectors)))

输出:

['A', 'C']

它按位置配对:

(di,si)di当且仅当 bool(si)=True(d_i,s_i) \mapsto d_i \quad \text{当且仅当 } bool(s_i)=True

并在 dataselectors 任意一个耗尽时停止。(docs.python.org)

这适合把“数据”和“选择结果”分开计算:

from itertools import compress

records = ["alice", "bob", "carol", "dave"]
valid = [True, False, True, False]

selected = compress(records, valid)
print(list(selected))

输出:

['alice', 'carol']

如果选择器来自无限迭代器,而数据有限,结果仍会在数据耗尽时结束;反过来,如果数据无限而选择器有限,结果会在选择器耗尽时结束。


8.2 filterfalse():保留谓词为假的元素

from itertools import filterfalse

values = [1, 2, 3, 4, 5]

print(list(filterfalse(lambda x: x % 2 == 0, values)))

输出:

[1, 3, 5]

传入 None 时,等价于按 bool 判断,保留假值:

from itertools import filterfalse

values = [0, 1, "", "python", None, [], [1]]

print(list(filterfalse(None, values)))

输出:

[0, '', None, []]

这里的“假”是 Python 布尔上下文中的假,不是字符串意义上的“无效”。如果业务规则是“状态必须等于 inactive”,应显式写出比较条件,避免把空字符串、零和 None 混为一谈。


9. 累积与相邻关系:accumulate()pairwise()

9.1 accumulate():输出每一步的中间状态

accumulate() 默认做加法:

from itertools import accumulate

print(list(accumulate([1, 2, 3, 4])))

输出:

[1, 3, 6, 10]

形式化地说,输入为 x1,x2,,xnx_1,x_2,\ldots,x_n,默认累积结果为:

y1=x1y_1=x_1

yi=yi1+xiy_i=y_{i-1}+x_i

如果提供二元函数 ff,则:

y1=x1y_1=x_1

yi=f(yi1,xi)y_i=f(y_{i-1},x_i)

例如计算运行最大值:

from itertools import accumulate

values = [3, 1, 5, 2, 4]

print(list(accumulate(values, max)))

输出:

[3, 3, 5, 5, 5]

9.2 initial 改变输出长度

from itertools import accumulate

print(list(accumulate([1, 2, 3], initial=100)))

输出:

[100, 101, 103, 106]

没有 initial 时,输出数量等于输入数量;有 initial 时,输出数量多一个:

输出长度=输入长度+1\text{输出长度}=\text{输入长度}+1

这使它适合表示“初始余额 → 每次更新后的余额”:

from itertools import accumulate

def update(balance, change):
    return balance + change

changes = [100, -30, 50, -20]

print(list(accumulate(changes, update, initial=0)))

输出:

[0, 100, 70, 120, 100]

注意 accumulate() 返回的是所有中间状态,而 functools.reduce() 只返回最终状态。二者都执行归约关系,但观测点不同:前者暴露状态序列,后者只保留最后结果。itertools 官方文档也将 accumulate()functools.reduce() 作为相关工具区分说明。(docs.python.org)


9.3 pairwise():相邻元素的重叠窗口

from itertools import pairwise

print(list(pairwise([10, 15, 13, 20])))

输出:

[(10, 15), (15, 13), (13, 20)]

输入长度为 nn 时,输出长度为:

max(n1,0)\max(n-1,0)

它适合计算相邻差值:

from itertools import pairwise

values = [10, 15, 13, 20]

differences = (
    current - previous
    for previous, current in pairwise(values)
)

print(list(differences))

输出:

[5, -2, 7]

pairwise() 只需要保留前一个元素,因此辅助空间是常量级。输入少于两个元素时,输出为空。(docs.python.org)


10. groupby():连续分组,而不是全局聚合

10.1 groupby() 按相邻键变化切组

from itertools import groupby

data = ["A", "A", "B", "B", "A"]

for key, group in groupby(data):
    print(key, list(group))

输出:

A ['A', 'A']
B ['B', 'B']
A ['A']

最后的 A 不会与最开始的 A 合并。groupby() 的分组条件是:

key(xi)=key(xi1)key(x_i)=key(x_{i-1})

只要键发生变化,就结束当前组并开始新组。

这与 SQL 的 GROUP BY 不同。SQL 通常会把所有相同键的记录聚合到一起;itertools.groupby() 只识别连续区间。官方文档明确指出,输入通常需要先按相同键函数排序,并特别强调它与 SQL GROUP BY 的行为差异。(docs.python.org)


10.2 正确用法:先排序,再分组

from itertools import groupby

records = [
    {"user": "alice", "score": 10},
    {"user": "bob", "score": 20},
    {"user": "alice", "score": 15},
]

records.sort(key=lambda record: record["user"])

for user, group in groupby(records, key=lambda record: record["user"]):
    rows = list(group)
    print(user, rows)

输出:

alice [{'user': 'alice', 'score': 10}, {'user': 'alice', 'score': 15}]
bob [{'user': 'bob', 'score': 20}]

排序的目的不是让 groupby() 更快,而是建立它所需要的输入不变量:

i,j,i<j<kkey(xi)=key(xk)\forall i,j,\quad i<j<k \land key(x_i)=key(x_k)

若希望相同键的元素属于同一个组,则中间不能出现其他键。排序可以保证相同键形成连续区间。


10.3 分组迭代器与主迭代器共享输入

这是 groupby() 最重要的生命周期规则:

from itertools import groupby

data = ["A", "A", "B", "B"]
groups = groupby(data)

key_a, group_a = next(groups)
key_b, group_b = next(groups)

print(key_a)          # A
print(list(group_a))  # []
print(key_b)          # B
print(list(group_b))  # ['B', 'B']

当主 groupby 迭代器推进到下一个组时,前一个分组已经被跳过。原因是:

  • 主迭代器和分组迭代器共享同一个底层输入;
  • 为了找到下一个键,主迭代器必须继续读取输入;
  • 一旦越过前一组,前一组的迭代器就无法回头。

如果后续还需要组内数据,必须立即物化:

from itertools import groupby

data = ["A", "A", "B", "B"]

saved_groups = [
    (key, list(group))
    for key, group in groupby(data)
]

print(saved_groups)

输出:

[('A', ['A', 'A']), ('B', ['B', 'B'])]

这会把所有组保存到内存中,内存复杂度由输入规模决定。若输入很大,应在每个组上立即完成处理,而不是把所有组转换成列表。


11. tee():分叉迭代器,但缓存速度差

11.1 为什么普通迭代器不能直接多次消费

source = iter([1, 2, 3])

first = list(source)
second = list(source)

print(first)   # [1, 2, 3]
print(second)  # []

如果两个消费者都需要看到完整输入,可以使用 tee()

from itertools import tee

source = iter([1, 2, 3])
left, right = tee(source)

print(list(left))   # [1, 2, 3]
print(list(right))  # [1, 2, 3]

tee() 返回多个独立迭代器,但它们并不复制源数据两份。更接近下面的数据流:

源迭代器
   ↓
共享缓存
  ↙   ↘
left  right

当两个分支同步消费时,缓存不会无限增长:

from itertools import tee

left, right = tee(range(5))

for a, b in zip(left, right):
    print(a, b)

输出:

0 0
1 1
2 2
3 3
4 4

但如果一个分支明显领先:

from itertools import tee

fast, slow = tee(range(1_000_000))

for _ in range(900_000):
    next(fast)

# 此时 slow 尚未消费的元素需要被缓存,以便 slow 未来读取

缓存规模大致与两个分支的消费进度差有关:

缓存量max(分支消费位置)min(分支消费位置)\text{缓存量} \approx \max(\text{分支消费位置})- \min(\text{分支消费位置})

因此,tee() 的内存风险不是分支数量本身,而是分支之间的最大进度差

官方文档还说明,由 tee() 返回的迭代器通常不应在多个线程中并发使用;即使底层输入本身支持并发,也不能据此推导出 tee() 分支安全。(docs.python.org)


11.2 不要把 tee() 当成廉价复制

如果输入是可重复遍历的容器,直接创建两个迭代器通常更清晰:

data = [1, 2, 3]

first = iter(data)
second = iter(data)

只有在源本身是一次性迭代器,例如文件流、生成器或网络读取器时,tee() 才解决“多个消费者都需要读取源”的问题。

但如果两个消费者处理速度差异很大,通常更适合:

  • 让一个消费者完成读取并显式写入队列;
  • 使用有界队列施加背压;
  • 将数据落盘后由多个消费者独立读取;
  • 或重新设计为单次遍历同时完成两类处理。

tee() 解决的是迭代语义,不是通用消息队列,也不提供持久化、确认、重试和故障恢复。


12. zip()map() 的截断语义

虽然 zip()map() 不属于 itertools,但它们经常与 itertools 组合,必须理解其输入边界。

from itertools import count, islice

values = ["a", "b", "c"]
numbers = count(1)

result = zip(values, numbers)

print(list(result))

输出:

[('a', 1), ('b', 2), ('c', 3)]

有限输入 values 使整个 zip() 结束,即使 numbers 无限。

相反:

from itertools import count

result = map(lambda x: x * 2, count())

这个结果也是无限迭代器,必须截断:

from itertools import count, islice

result = islice(map(lambda x: x * 2, count()), 5)

print(list(result))

输出:

[0, 2, 4, 6, 8]

组合无限流时,应明确回答两个问题:

  1. 哪个组件负责终止?
  2. 终止之前最多会消费多少上游数据?

如果这两个问题没有答案,流水线很可能在测试中正常、在生产中永不结束。


13. 一个完整的惰性流水线

下面构造一个订单日志处理流程:

  1. 生成订单编号;
  2. 使用 repeat() 添加固定来源;
  3. 使用 starmap() 构造记录;
  4. 使用 filterfalse() 丢弃无效记录;
  5. 使用 batched() 分批;
  6. 使用 accumulate() 观察累计金额。
from itertools import (
    accumulate,
    batched,
    filterfalse,
    repeat,
    starmap,
)

def build_record(order_id, source):
    amount = order_id * 10
    return {
        "order_id": order_id,
        "source": source,
        "amount": amount,
        "valid": order_id % 3 != 0,
    }

order_ids = range(1, 8)

records = starmap(
    build_record,
    zip(order_ids, repeat("api")),
)

valid_records = filterfalse(
    lambda record: not record["valid"],
    records,
)

for batch in batched(valid_records, 2):
    amounts = (record["amount"] for record in batch)

    print("batch:", batch)
    print("running total:", list(accumulate(amounts)))

输出:

batch: (
    {'order_id': 1, 'source': 'api', 'amount': 10, 'valid': True},
    {'order_id': 2, 'source': 'api', 'amount': 20, 'valid': True}
)
running total: [10, 30]

batch: (
    {'order_id': 4, 'source': 'api', 'amount': 40, 'valid': True},
    {'order_id': 5, 'source': 'api', 'amount': 50, 'valid': True}
)
running total: [40, 90]

batch: (
    {'order_id': 7, 'source': 'api', 'amount': 70, 'valid': True},
)
running total: [70]

数据流可以表示为:

flowchart LR
    A[order_ids] --> B[zip]
    C[repeat api] --> B
    B --> D[starmap build_record]
    D --> E[filterfalse]
    E --> F[batched size=2]
    F --> G[batch 内 accumulate]

每个阶段的内存边界不同:

  • rangerepeatzipstarmapfilterfalse:主要保存迭代器和函数状态;
  • batched(..., 2):每次最多保存两个记录;
  • list(accumulate(...)):只保存当前批次的累计结果;
  • batch 本身是元组,因此当前批次的记录会暂时保留。

如果将最终结果写入文件,而不是调用 list(),整个输入就不需要一次性驻留内存。


14. 常见误解与失败表现

14.1 误解:list(itertools_function(...)) 只是查看结果

result = list(combinations(range(100_000), 2))

这会尝试保存全部组合,数量接近 50 亿。即使组合生成器本身是惰性的,list() 也会主动取消这种惰性,要求完整物化结果。

诊断方法是先估算结果数量:

from math import comb

n = 100_000
r = 2

print(comb(n, r))

如果只需要检查前几个结果:

from itertools import combinations, islice

preview = islice(combinations(range(100_000), 2), 10)
print(list(preview))

14.2 误解:groupby() 会把所有相同键合并

from itertools import groupby

data = ["a", "b", "a"]

print([(key, list(group)) for key, group in groupby(data)])

输出:

[('a', ['a']), ('b', ['b']), ('a', ['a'])]

如果业务需要全局聚合,应使用字典:

from collections import defaultdict

groups = defaultdict(list)

for value in data:
    groups[value].append(value)

print(dict(groups))

groupby() 适合已经按键排序的数据流、日志分段和连续状态区间;字典聚合适合无序输入的全局归并。二者不是等价替代品。


14.3 误解:tee() 会复制出多个独立数据集

from itertools import tee

source = (value for value in range(10))
first, second = tee(source)

这不会立即复制 10 个元素。元素会在某个分支请求时从源读取,并为尚未追上的分支缓存。

如果 first 被完全消费,而 second 长时间不消费,缓存会持续增长。诊断时应记录:

  • 每个分支已经消费的位置;
  • 源迭代器的总产出;
  • 分支之间的滞后距离;
  • 进程内存是否随滞后距离增长。

14.4 误解:所有无限工具都只占常量内存

count() 和有限次数的 repeat() 通常不需要保存历史输出;cycle() 则必须保存第一轮输入。

from itertools import cycle, islice

result = islice(cycle(range(1_000_000)), 10)

print(list(result))

虽然只输出 10 个元素,cycle() 为了支持未来循环,可能已经读取并保存整个有限输入。这是“输出数量小”和“中间缓存小”之间的区别。


14.5 误解:islice() 会像列表切片一样随机访问

列表切片可以直接根据索引定位;迭代器没有随机访问能力。下面的操作:

from itertools import islice

result = islice(source, 1_000_000, 1_000_001)

需要从 source 依次消费前 1,000,000 个元素,不能跳过而不读取。

因此,islice() 适合顺序流的前缀、窗口和抽样,不适合把生成器伪装成支持高效随机访问的数据结构。


15. 选择工具时应先判断三个边界

15.1 输入是否有限

以下工具需要有限输入才能完成预期工作:

  • product()
  • permutations()
  • combinations()
  • combinations_with_replacement()
  • cycle() 的第一轮输入。

其中 product() 会先完全消费输入;cycle() 对有限输入会缓存全部元素。(docs.python.org)


15.2 结果空间是否可接受

结果数量的增长关系必须先算清楚:

工具 结果数量
product() ni\prod n_i
product(A, repeat=r) nrn^r
permutations(n, r) n!(nr)!\frac{n!}{(n-r)!}
combinations(n, r) n!r!(nr)!\frac{n!}{r!(n-r)!}
combinations_with_replacement(n, r) (n+r1)!r!(n1)!\frac{(n+r-1)!}{r!(n-1)!}

生成器只降低“结果同时驻留内存”的压力,不降低必须枚举的结果总数。


15.3 是否存在共享状态或缓存

可以把常见工具分成三种内存模式:

严格流式

通常只保存当前元素和少量状态:

  • chain()
  • compress()
  • filterfalse()
  • islice()
  • pairwise()
  • starmap()
  • takewhile()
  • dropwhile()
  • accumulate()
  • batched() 的当前批次。

输入池化

会先保存输入元素:

  • product()
  • permutations()
  • combinations()
  • combinations_with_replacement()

历史缓存或共享缓存

可能随着输入长度或分支差距增长:

  • cycle():保存第一轮输入;
  • tee():保存领先分支已经消费的数据;
  • groupby():分组迭代器共享底层输入,但本身不会自动保存历史分组。

这个分类比“是否惰性”更能预测生产环境中的内存行为。


16. 版本边界:Python 3.14 中应注意什么

本文使用的是 Python 3.14 标准库语义。几个与版本相关的 API:

  • pairwise():Python 3.10 加入;
  • batched():Python 3.12 加入;
  • batched(..., strict=True):Python 3.13 加入;
  • accumulate(..., initial=...):Python 3.8 加入。

因此,如果代码需要运行在更早版本上,不能只根据当前解释器的文档编写。特别是 batched() 在早期版本中不存在,strict 参数也不是所有版本都支持。上述版本信息和参数变化以 Python 3.14 官方文档为准。(docs.python.org)


结语:惰性是数据流策略,边界才是资源模型

itertools 的核心价值不在于把循环压缩成一行,而在于把计算组织成可组合的数据流:

  • count()repeat() 提供可控或无限的输入;
  • islice() 为无限流建立终止边界;
  • chain()batched()zip_longest() 改变数据形状;
  • filterfalse()compress()takewhile()dropwhile() 控制选择;
  • accumulate()pairwise() 暴露局部状态;
  • groupby() 根据连续键切分输入;
  • product()permutations() 和组合工具枚举候选空间;
  • tee() 在一次性输入上建立多个消费视图。

真正需要掌握的是每个工具的状态和边界:

  1. 它什么时候消费上游?
  2. 它是否会预读?
  3. 它会保存哪些历史数据?
  4. 输入无限时是否仍能终止?
  5. 输出数量是线性、乘法级、排列级还是组合级?
  6. 下游是否会通过 list()sorted() 或缓存重新物化整个流?

能回答这些问题,才算真正理解“惰性组合、无限序列、分组和内存边界”。


系列导航与关联阅读

官方资料

本文依据 Python 官方文档、相关 PEP 与生态项目官方文档重新梳理;正文、示例与工程清单由 WR BLOG 编写。