python | Python中的itertools模块使用技巧

本文来源公众号“python”，仅用于学术分享，侵权删，干货满满。

原文链接：Python中的itertools模块使用技巧

itertools模块是Python标准库中的一个高效工具箱，提供了许多用于处理迭代器的函数。通过itertools模块，可以简化复杂的循环和迭代操作，大幅提升代码效率。itertools模块内的函数涵盖了排列组合、无限迭代器、分组、过滤等丰富功能，特别适合数据处理、序列操作等场景。

1 无限迭代器

itertools模块提供了几种常用的无限迭代器，适用于需要生成无限序列的场景。常见的无限迭代器包括count、cycle和repeat。

1.1 count：生成等差数列

count(start, step)函数从start开始生成无限等差数列，每次递增step。

以下示例展示了如何生成一个从10开始，每次增加5的无限数列：

import itertools

for num in itertools.count(10, 5):
    if num > 30:
        break
    print(num)
# 输出：10, 15, 20, 25, 30

1.2 cycle：循环遍历序列

cycle(iterable)将对序列进行无限循环遍历，适用于需要反复访问的场景。

以下示例展示了如何无限循环一个列表的元素：

colors = ['red', 'green', 'blue']
cycled_colors = itertools.cycle(colors)

for _ in range(6):
    print(next(cycled_colors))
# 输出：red, green, blue, red, green, blue

1.3 repeat：重复生成指定元素

repeat(item, times)生成一个重复的元素序列，重复次数为times，若不指定times则会无限重复。

以下示例展示了生成5次相同元素的序列：

for item in itertools.repeat("hello", 5):
    print(item)
# 输出：hello, hello, hello, hello, hello

2 排列与组合

itertools模块提供了几种函数，用于生成序列的排列和组合，常用于数据处理和分析中。

2.1 permutations：生成排列

permutations(iterable, r)生成长度为r的排列，若不指定r则默认生成与输入序列等长的排列。

以下示例展示了列表中元素的全排列：

data = [1, 2, 3]
for perm in itertools.permutations(data):
    print(perm)
# 输出：(1, 2, 3), (1, 3, 2), (2, 1, 3), (2, 3, 1), (3, 1, 2), (3, 2, 1)

2.2 combinations：生成组合

combinations(iterable, r)生成长度为r的组合，不考虑元素顺序。

以下示例展示了从列表中选择2个元素的所有组合：

for comb in itertools.combinations(data, 2):
    print(comb)
# 输出：(1, 2), (1, 3), (2, 3)

2.3 combinations_with_replacement：生成有重复的组合

combinations_with_replacement(iterable, r)允许元素重复生成组合。

以下示例展示了带重复的组合：

for comb in itertools.combinations_with_replacement(data, 2):
    print(comb)
# 输出：(1, 1), (1, 2), (1, 3), (2, 2), (2, 3), (3, 3)

3 累积与分组

itertools还提供了数据累积与分组的函数，适用于需要进行累计操作或分组的场景。

3.1 accumulate：累积求和

accumulate(iterable, func)生成累计结果序列，默认执行加法，可以通过func参数指定其他运算。

以下示例展示了累积求和：

data = [1, 2, 3, 4]
for acc in itertools.accumulate(data):
    print(acc)
# 输出：1, 3, 6, 10

3.2 groupby：分组操作

groupby(iterable, key)根据key分组，key是一个函数，用于定义分组依据。

以下示例展示了将字符串按相同字符分组：

data = "aaabbc"
for key, group in itertools.groupby(data):
    print(key, list(group))
# 输出：a ['a', 'a', 'a'], b ['b', 'b'], c ['c']

注意，groupby只能对相邻元素进行分组，因此在使用前需对数据进行排序。

4 组合迭代器

itertools模块提供了一些组合迭代器，用于将多个迭代器组合在一起，常见的组合迭代器包括chain、zip_longest和product。

4.1 chain：连接多个迭代器

chain(*iterables)用于将多个迭代器连接成一个迭代器，适合用于连接多个序列。

以下示例展示了将两个列表连接成一个迭代器：

for item in itertools.chain([1, 2, 3], ['a', 'b', 'c']):
    print(item)
# 输出：1, 2, 3, a, b, c

4.2 zip_longest：填充组合

zip_longest(*iterables, fillvalue=None)用于将长度不一的迭代器进行组合，较短的迭代器用fillvalue填充。

以下示例展示了将不同长度的列表组合成一对：

list1 = [1, 2]
list2 = ['a', 'b', 'c']
for item in itertools.zip_longest(list1, list2, fillvalue="X"):
    print(item)
# 输出：(1, 'a'), (2, 'b'), ('X', 'c')

4.3 product：笛卡尔积

product(*iterables, repeat=1)用于生成多个迭代器的笛卡尔积，适用于生成所有可能的排列组合。

以下示例展示了两个列表的笛卡尔积：

for item in itertools.product([1, 2], ['a', 'b']):
    print(item)
# 输出：(1, 'a'), (1, 'b'), (2, 'a'), (2, 'b')

5 高级技巧：多种方法结合应用

结合使用itertools的多种方法可以实现一些复杂的数据处理操作。

5.1 滑动窗口

滑动窗口是一种遍历序列的方式，可用于序列分析等操作。

可以结合zip和itertools.islice实现滑动窗口：

from itertools import islice

def sliding_window(iterable, n):
    iters = [islice(iterable, i, None) for i in range(n)]
    return zip(*iters)

data = [1, 2, 3, 4, 5]
for window in sliding_window(data, 3):
    print(window)
# 输出：(1, 2, 3), (2, 3, 4), (3, 4, 5)

5.2 批量处理

在处理大数据时，可以将数据分批处理。

结合islice实现批量数据处理：

def batched(iterable, n):
    it = iter(iterable)
    while True:
        batch = list(islice(it, n))
        if not batch:
            break
        yield batch

data = range(1, 11)
for batch in batched(data, 3):
    print(batch)
# 输出：[1, 2, 3], [4, 5, 6], [7, 8, 9], [10]

5.3 嵌套排列组合

在数据分析中，经常需要嵌套排列组合，例如生成多种分类的所有组合，可以通过嵌套product实现：

colors = ['red', 'green']
sizes = ['S', 'M', 'L']
styles = ['casual', 'formal']

for combo in itertools.product(colors, sizes, styles):
    print(combo)
# 输出所有颜色、尺寸和样式的组合

6 总结

itertools模块是Python处理迭代器的强大工具箱，提供了丰富的函数来简化复杂的迭代和数据操作。通过count、cycle等无限迭代器，可以轻松生成无限序列；permutations和combinations使得排列和组合的实现更加便捷；accumulate与groupby则提供了数据累积与分组功能。此外，chain、zip_longest和product等组合迭代器支持多序列的灵活操作。无论是数据分析、序列操作，还是批量处理和滑动窗口，itertools都能提供高效的解决方案。这些工具函数不仅能简化代码，还能提升运行效率。

THE END !

文章结束，感谢阅读。您的点赞，收藏，评论是我继续更新的动力。大家有推荐的公众号可以评论区留言，共同学习，一起进步。