python进阶系列 - 07 itertools 迭代器

在Python的标准库当中有这么一个神奇的工具库，它能让你使用最简单的方式写出更简洁高效的代码，这就是itertools，使用这个工具能为你生成一个优雅的迭代器。

这个模块提供了一些工具来处理迭代器。 简单地说，迭代器是一种可以在 for 循环中使用的数据类型。 Python 中最常见的迭代器是列表。
可以在官网获得itertools的全部用法。

product() 笛卡尔积

这个工具计算输入迭代器的笛卡尔积。 它等价于嵌套的 for 循环。例如，product(A, B) 返回((x,y) for x in A for y in B) 相同。

代码：

from itertools import product

prod = product([1, 2], [3, 4])
# 注意，我们将迭代器转换为列表以便打印
print(list(prod))
# 以允许迭代器的笛卡尔积，请指定重复次数
prod = product([1, 2], [3], repeat=2)
print(list(prod))

结果：

[(1, 3), (1, 4), (2, 3), (2, 4)]
[(1, 3, 1, 3), (1, 3, 2, 3), (2, 3, 1, 3), (2, 3, 2, 3)]

permutations() 排列

该函数返回迭代器中元素的所有可能排列（无重复元素）。

代码：

from itertools import permutations

perm = permutations([1, 2, 3])
print(list(perm))

# 可选参数：排列的长度。
perm = permutations([1, 2, 3], 2)
print(list(perm))

结果：

[(1, 2, 3), (1, 3, 2), (2, 1, 3), (2, 3, 1), (3, 1, 2), (3, 2, 1)]
[(1, 2), (1, 3), (2, 1), (2, 3), (3, 1), (3, 2)]

combinations() 组合

组合

如果输入迭代器是排序的，组合元组将被生成在排序的顺序中。
combinations() 不允许重复元素，combinations_with_replacement() 允许重复元素。

代码：

from itertools import combinations, combinations_with_replacement

# 1. 第二个参数是必需的，它指定了输出元组的长度。
# 2. 输出所有2个元素的组合。无重复元素。
# 3. 注意是组合，不是排列，也就是[1, 2]和[2, 1]不会被同时输出。
comb = combinations([1, 2, 3, 4], 2)
print(list(comb))

comb = combinations_with_replacement([1, 2, 3, 4], 2)  # 输出所有2个元素的组合，包括重复元素。
print(list(comb))

[(1, 2), (1, 3), (1, 4), (2, 3), (2, 4), (3, 4)]
[(1, 1), (1, 2), (1, 3), (1, 4), (2, 2), (2, 3), (2, 4), (3, 3), (3, 4), (4, 4)]

accumulate() 累加器

accumulate()函数基于给定的函数返回一个可迭代对象，将一系列归约值汇总在一起，遍历迭代器得出当前汇总值。默认的函数是operator.add()。

代码：

from itertools import accumulate

# 返回累计和
acc = accumulate([1, 2, 3, 4])
print(list(acc))

# 可以使用其他二元函数
import operator

acc = accumulate([1, 2, 3, 4], func=operator.mul)  # 返回累计乘
print(list(acc))

acc = accumulate([1, 5, 2, 6, 3, 4], func=max)  # 返回累计最大值
print(list(acc))

结果：

[1, 3, 6, 10]
[1, 2, 6, 24]
[1, 5, 5, 6, 6, 6]

groupby() 分组

Python 用groupby()切分迭代器，通过对每个元素应用key函数进行求值，groupby()函数将一个迭代器切分为多个小迭代器。如果后一个元素的key值等于前一个元素的key
值，会将这两个元素放在同一个分组中；如果与前一个元素的key值不同，则当前分组结束，将当前元素放到新的分组中。

groupby()函数的输出是一系列二元组，每个元组包括一个key值和包含该组元素的迭代器，该迭代器的元素可以通过转换保存为元组，也可以归约为汇总值。这种情况下将不会保留迭代器中的值。

代码：

from itertools import groupby


# 用函数作为键
def smaller_than_3(x):
    return x < 3


group_obj = groupby([1, 2, 3, 4], key=smaller_than_3)
for key, group in group_obj:
    print(key, list(group))

# 或者使用lambda表达式，例如：单词包含'i'：
group_obj = groupby(["hi", "nice", "hello", "cool"], key=lambda x: "i" in x)
for key, group in group_obj:
    print(key, list(group))
persons = [{'name': 'Tim', 'age': 25}, {'name': 'Dan', 'age': 25},
           {'name': 'Lisa', 'age': 27}, {'name': 'Claire', 'age': 28}]

# 返回一个字典，键是年龄，值是一个列表，列表中包含名字。
for key, group in groupby(persons, key=lambda x: x['age']):
    print(key, list(group))

结果：

True [1, 2]
False [3, 4]
True ['hi', 'nice']
False ['hello', 'cool']
25 [{'name': 'Tim', 'age': 25}, {'name': 'Dan', 'age': 25}]
27 [{'name': 'Lisa', 'age': 27}]
28 [{'name': 'Claire', 'age': 28}]

无限迭代器: count(), cycle(), repeat()

python有3个无限迭代，分别为count(),cycle(),repeat()

• count(start[, step]) 无限计数递增迭代，参数start起始数据，可选参数step为步长
• cycle(p) 无限循环周期性迭代，参数p，周期性循环p序列的数据
• repeat(elem [,n]) 重复迭代，对elem重复迭代n次

代码:

from itertools import count, cycle, repeat

# count(x)：从x开始，每次增加1
for i in count(10):
    print(i)
    if i >= 13:
        break

# cycle(iterable)：无限循环迭代器
print("")
sum = 0
for i in cycle([1, 2, 3]):
    print(i)
    sum += i
    if sum >= 12:
        break

# repeat(x, n)：重复x n次
print("")
for i in repeat("A", 3):
    print(i)

结果:

10
11
12
13

1
2
3
1
2
3

A
A
A

小节

如文所示，itertools模块，提供便利的迭代器工具，方便快速完成如排列、组合、重复等工作中高频常见的任务。

如果你觉得这篇文章对你有帮助，欢迎大家点赞、收藏、支持！。

pythontip 出品，Happy Coding！

公众号: 夸克编程