如何将迭代器（iterator）均匀划分为N个块（chunk）？

需求：写一个split_n函数，把任意有限迭代器分成n个迭代器类型的连续块，要求：

1. 绝对不能把原迭代器转成列表（防止超大迭代器耗尽内存）
2. 块与块之间的元素数量尽可能均等，无法整除时，前几个块多1个元素

预期效果示例 #

# 示例1：整除场景
iter1 = iter(range(10))
chunks = split_n(iter1, n=2)
# 遍历块（转列表仅为展示）
for chunk in chunks:
    print(list(chunk))
# 输出：
# [0, 1, 2, 3, 4]
# [5, 6, 7, 8, 9]

# 示例2：整除场景
iter2 = iter(range(20))
chunks = split_n(iter2, n=4)
for chunk in chunks:
    print(list(chunk))
# 输出：
# [0, 1, 2, 3, 4]
# [5, 6, 7, 8, 9]
# [10, 11, 12, 13, 14]
# [15, 16, 17, 18, 19]

# 示例3：无法整除场景
iter3 = iter(range(30))
chunks = split_n(iter3, n=4)
for chunk in chunks:
    print(list(chunk))
# 输出：
# [0, 1, 2, 3, 4, 5, 6, 7]
# [8, 9, 10, 11, 12, 13, 14, 15]
# [16, 17, 18, 19, 20, 21, 22]
# [23, 24, 25, 26, 27, 28, 29]

实现代码 #

from itertools import islice, tee

def split_n(iterator, n):
    # 复制迭代器，一个用于计数，一个用于分块
    count_iter, chunk_iter = tee(iterator, 2)
    
    # 统计总元素数，仅遍历一次，无内存积累
    total = sum(1 for _ in count_iter)
    
    base_size, remainder = divmod(total, n)
    
    # 生成每个块的迭代器
    for i in range(n):
        current_size = base_size + 1 if i < remainder else base_size
        if current_size <= 0:
            break
        yield islice(chunk_iter, current_size)

代码逻辑说明 #

• 内存控制：用tee复制迭代器，计数环节用sum(1 for _ in ...)遍历迭代器，不会存储任何元素，仅做计数操作，内存开销为O(1)，完全符合不转列表的要求。
• 均匀分配：通过divmod计算基础块长度和剩余元素数，前remainder个块各多1个元素，确保所有块的长度差不超过1，达到最均等的分配效果。
• 迭代器特性保留：每个块由islice生成，本质是迭代器，只有在遍历块时才会从原迭代器中取元素，全程不占用额外内存存储元素。

注意事项 #

• 原迭代器是一次性的，调用split_n后原迭代器会被耗尽，无法再次使用。
• 该函数仅适用于有限长度的迭代器，传入无限迭代器会导致计数环节无限循环。

内容的提问来源于stack exchange，提问作者Thegerdfather