【python】常用内置模块itertools的用法（迭代器)

---

1. 简介

python内置模块itertools中封装了很多高效的迭代器，本文将讲述10个最为常用的迭代器函数。

2. 无限迭代器

2.1. count(start, [step])

输出自start参数开始的无限长度的自然数序列，step参数设置步长，可选。

>>> import itertools as its
>>> a = its.count(1, 2)
>>> for i in a:
...     print(i)
...
1
3
5
...
使用快捷键Ctrl+C暂停打印

2.2. cycle( p)

将一个序列无限循环下去，p参数指定待循环的序列。

>>> import itertools as its
>>> a = its.cycle('ABC')
>>> for i in a:        
...     print(i)       
... 
A
B
C
A
B
C
...
使用快捷键Ctrl+C暂停打印

2.3. repeat(elem, [n])

把一个元素重复n次，如果不指定参数n，则默认无限次。

>>> import itertools as its
>>> a = its.repeat(10, 3)
>>> for i in a:
...     print(i)        
... 
10  
10  
10 

3. 通用迭代器

3.1. accumulate(p, [func])

依次累加序列p中的元素。

>>> import itertools as its
>>> a = its.accumulate([1, 2, 3, 4, 5])
>>> for i in a:
...     print(i)
... 
1   # 1
3   # 1+2
6   # 1+2+3
10  # 1+2+3+4
15  # 1+2+3+4+5

3.2. chain(p, q, […])

串联多个迭代器，形成一个更大的迭代器。

>>> import itertools as its
>>> a = its.chain('ABD', 'DEF')
>>> for i in a:
...     print(i)
... 
A   
B   
D   
D   
E   
F 

3.3. takewhile(pred, seq)

迭代选择器。

>>> import itertools as its
>>> a = its.takewhile(lambda x: x < 5, [1, 3, 5, 7])
>>> for i in a:
...     print(i)
... 
1   
3  

4. 组合迭代器

4.1. product(p, q, […], [repeat=1])

笛卡尔积，等效于嵌套的for循环：

>>> import itertools as its
>>> a = its.product('ABC', repeat=2)
>>> for i in a:
...     print(i)
... 
('A', 'A')
('A', 'B')
('A', 'C')
('B', 'A')
('B', 'B')
('B', 'C')
('C', 'A')
('C', 'B')
('C', 'C')

等价于：

>>> for x in 'ABC':
...     for y in 'ABC':
...         print((x, y))
... 
('A', 'A')
('A', 'B')
('A', 'C')
('B', 'A')
('B', 'B')
('B', 'C')
('C', 'A')
('C', 'B')
('C', 'C')

提问：既然可以用嵌套的for循环来代替，为什么还要封装该函数？

解答：一两个for循环可以，但如果你需要写七八个for循环，代码就很难看了，所以一般都会用该函数替换嵌套的for循环。

该函数可以应用到密码字典的生成中，比如生成一个包含所有由数字组成的长度为8的密码字典：

import itertools as its
iterator = its.product('0123456789', repeat=8)
with open('pwd.txt', 'w') as f:
    for i in iterator:
        f.write(''.join(i)+'\n')

4.2. permutations(p, [r])

同product()函数，但元素不允许重复，不会生成AA这样的：

>>> import itertools as its
>>> a = its.permutations('ABC', r=2)
>>> for i in a:
...     print(i)
... 
('A', 'B')
('A', 'C')
('B', 'A')
('B', 'C')
('C', 'A')
('C', 'B')

4.3. combinations(p, r)

同permutations()函数，但AB和BA这样的视为同一元素，只取AB：

>>> import itertools as its
>>> a = its.combinations('ABC', r=2)
>>> for i in a:
...     print(i)
...
('A', 'B')
('A', 'C')
('B', 'C')

4.4. combinations_with_replacement(p, r)

同product()函数，允许重复，但AB和BA这样的视为同一元素，只取AB：

>>> import itertools as its
>>> a = its.combinations_with_replacement('ABC', r=2)
>>> for i in a:
...     print(i)
...
('A', 'A')
('A', 'B')
('A', 'C')
('B', 'B')
('B', 'C')
('C', 'C')

5. 引用

https://docs.python.org/3.10/library/itertools.html#itertools-recipes