python进阶系列 - 07 itertools 迭代器

169 阅读3分钟

在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!

公众号: 夸克编程