高效的itertools模块

发表于 | 分类于 | | 阅读次数: |
字数统计: 2,449 字 | 阅读时长 ≈ 12 分钟

我们知道，迭代器的特点是：惰性求值（Lazy evaluation），即只有当迭代至某个值时，它才会被计算，这个特点使得迭代器特别适合于遍历大文件或无限集合等，因为我们不用一次性将它们存储在内存中。

Python 内置的 itertools 模块包含了一系列用来产生不同类型迭代器的函数或类，这些函数的返回都是一个迭代器，我们可以通过 for 循环来遍历取值，也可以使用 next() 来取值。

itertools 模块提供的迭代器函数有以下几种类型：

  • 无限迭代器：生成一个无限序列，比如自然数序列 1, 2, 3, 4, ...；
  • 有限迭代器：接收一个或多个序列（sequence）作为参数，进行组合、分组和过滤等；
  • 组合生成器：序列的排列、组合，求序列的笛卡儿积等；

无限迭代器

itertools 模块提供了三个函数（事实上，它们是类）用于生成一个无限序列迭代器：

  • count(firstval=0, step=1)

    创建一个从 firstval (默认值为 0) 开始，以 step (默认值为 1) 为步长的的无限整数迭代器

  • cycle(iterable)

    对 iterable 中的元素反复执行循环，返回迭代器

  • repeat(object [,times]

    反复生成 object，如果给定 times，则重复次数为 times，否则为无限

下面，让我们看看一些例子。

count

count() 接收两个参数，第一个参数指定开始值，默认为 0，第二个参数指定步长，默认为 1：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
In [1]: import itertools

In [2]: nums = itertools.count()

In [3]: for i in nums:
   ...:     if i > 6:
   ...:         break
   ...:     print(i)
   ...:
0
1
2
3
4
5
6
1
2
3
4
5
6
7
8
9
10
11
12
13
14
In [4]: nums = itertools.count(10, 2) # 指定开始值和步长

In [5]: for i in nums:
   ...:     if i > 20:
   ...:         break
   ...:     print(i)
   ...:
   ...:
10
12
14
16
18
20
cycle

cycle() 用于对 iterable 中的元素反复执行循环：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
In [6]: cycle_strings = itertools.cycle('ABC')

In [7]: i = 1

In [8]: for string in cycle_strings:
   ...:     if i == 10:
   ...:         break
   ...:     print(i, string)
   ...:     i += 1
   ...:
1 A
2 B
3 C
4 A
5 B
6 C
7 A
8 B
9 C
repeat

repeat() 用于反复生成一个 object：

1
2
3
4
5
6
In [10]: for item in itertools.repeat('hello world', 3):
    ...:     print(item)
    ...:
hello world
hello world
hello world
1
2
3
4
5
6
In [11]: for item in itertools.repeat([1, 2, 3, 4], 3):
    ...:     print(item)
    ...:
[1, 2, 3, 4]
[1, 2, 3, 4]
[1, 2, 3, 4]

有限迭代器

itertools 模块提供了多个函数（类），接收一个或多个迭代对象作为参数，对它们进行组合、分组和过滤等：

  • chain()
  • compress()
  • dropwhile()
  • groupby()
  • filterfalse()
  • islice()
  • starmap()
  • tee()
  • takewhile()
  • zip_longest()
chain

chain 的使用形式如下：

1
chain(iterable1, iterable2, iterable3, ...)

chain 接收多个可迭代对象作为参数，将它们『连接』起来，作为一个新的迭代器返回。

1
2
3
4
5
6
7
8
9
10
11
In [12]: from itertools import chain

In [13]: for item in chain([1, 2, 3], ['a', 'b', 'c']):
    ...:     print(item)
    ...:
1
2
3
a
b
c

chain 还有一个常见的用法：

1
chain.from_iterable(iterable)

接收一个可迭代对象作为参数，返回一个迭代器：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
In [20]: from itertools import chain

In [21]: string = chain.from_iterable('ABCD')

In [22]: next(string)
Out[22]: 'A'

In [23]: next(string)
Out[23]: 'B'

In [24]: next(string)
Out[24]: 'C'

In [25]: next(string)
Out[25]: 'D'

In [26]: next(string)
---------------------------------------------------------------------------
StopIteration                             Traceback (most recent call last)
<ipython-input-26-50005732e4cf> in <module>()
----> 1 next(string)

StopIteration:
compress

compress 的使用形式如下：

1
compress(data, selectors)

compress 可用于对数据进行筛选，当 selectors 的某个元素为 true 时，则保留 data 对应位置的元素，否则去除：

1
2
3
4
5
6
7
8
9
10
In [27]: from itertools import compress

In [28]: list(compress('ABCDEF', [1, 1, 0, 1, 0, 1]))
Out[28]: ['A', 'B', 'D', 'F']

In [29]: list(compress('ABCDEF', [1, 1, 0, 1]))
Out[29]: ['A', 'B', 'D']

In [30]: list(compress('ABCDEF', [True, False, True]))
Out[30]: ['A', 'C']
dropwhile

dropwhile 的使用形式如下：

1
dropwhile(predicate, iterable)

其中，predicate 是函数，iterable 是可迭代对象。对于 iterable 中的元素，如果 predicate(item) 为 true，则丢弃该元素，否则返回该项及所有后续项。

1
2
3
4
5
6
7
In [31]: from itertools import dropwhile

In [32]: list(dropwhile(lambda x: x < 5, [1, 3, 6, 2, 1]))
Out[32]: [6, 2, 1]

In [33]: list(dropwhile(lambda x: x > 3, [2, 1, 6, 5, 4]))
Out[33]: [2, 1, 6, 5, 4]
groupby

groupby 用于对序列进行分组，它的使用形式如下:

1
groupby(iterable[, keyfunc])

其中，iterable 是一个可迭代对象，keyfunc 是分组函数，用于对 iterable 的连续项进行分组，如果不指定，则默认对 iterable 中的连续相同项进行分组，返回一个 (key, sub-iterator) 的迭代器。

1
2
3
4
5
6
7
8
9
10
In [34]: from itertools import groupby

In [35]: for key, value_iter in groupby('aaabbbaaccd'):
    ...:     print(key, ':', list(value_iter))
    ...:
a : ['a', 'a', 'a']
b : ['b', 'b', 'b']
a : ['a', 'a']
c : ['c', 'c']
d : ['d']
1
2
3
4
5
6
7
8
9
10
11
In [36]: data = ['a', 'bb', 'ccc', 'dd', 'eee', 'f']

In [37]: for key, value_iter in groupby(data, len):
    ...:     print(key, ':', list(value_iter))
    ...:
1 : ['a']
2 : ['bb']
3 : ['ccc']
2 : ['dd']
3 : ['eee']
1 : ['f']
1
2
3
4
5
6
7
8
9
10
In [38]: data = ['a', 'bb', 'cc', 'ddd', 'eee', 'f']

In [40]: for key, value_iter in groupby(data, len):
    ...:     print(key, ':', list(value_iter))
    ...:
    ...:
1 : ['a']
2 : ['bb', 'cc']
3 : ['ddd', 'eee']
1 : ['f']
filterfalse

filterfalse 的使用形式如下：

1
filterfalse(function or None, sequence)

它将 iterable 中 function(item) 为 False 的元素组成一个迭代器返回，如果 function 是 None，则返回 iterable 中所有计算为 False 的项。

1
2
3
4
5
6
7
In [42]: from itertools import filterfalse

In [43]: list(filterfalse(lambda x: x < 6, range(10)))
Out[43]: [6, 7, 8, 9]

In [44]: list(filterfalse(None, [0, 1, 2, 0, 3, 4]))
Out[44]: [0, 0]
islice

islice 是切片选择，它的使用形式如下：

1
islice(iterable, [start,] stop [, step])

其中，iterable 是可迭代对象，start 是开始索引，stop 是结束索引，step 是步长，start 和 step 可选。

1
2
3
4
5
6
7
8
9
10
11
12
13
In [46]: from itertools import count, islice

In [47]: list(islice([10, 6, 2, 8, 1, 3, 9], 5))
Out[47]: [10, 6, 2, 8, 1]

In [48]: list(islice(count(), 6))
Out[48]: [0, 1, 2, 3, 4, 5]

In [49]: list(islice(count(), 3, 10))
Out[49]: [3, 4, 5, 6, 7, 8, 9]

In [50]: list(islice(count(), 3, 10 ,2))
Out[50]: [3, 5, 7, 9]
tee

tee 的使用形式如下：

1
tee(iterable [,n])

tee 用于从 iterable 创建 n 个独立的迭代器，以元组的形式返回，n 的默认值是 2。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
In [60]: from itertools import tee

In [61]: tee('abcd')
Out[61]: (<itertools._tee at 0x106288a48>, <itertools._tee at 0x106277248>)

In [62]: iter1, iter2 = tee('abcde')

In [63]: list(iter1)
Out[63]: ['a', 'b', 'c', 'd', 'e']

In [64]: list(iter2)
Out[64]: ['a', 'b', 'c', 'd', 'e']

In [65]: tee('abc', 3)
Out[65]:
(<itertools._tee at 0x10613e648>,
 <itertools._tee at 0x1062246c8>,
 <itertools._tee at 0x1062247c8>)
takewhile

takewhile 的使用形式如下：

1
takewhile(predicate, iterable)

其中，predicate 是函数，iterable 是可迭代对象。对于 iterable 中的元素，如果 predicate(item) 为 true，则保留该元素，只要 predicate(item) 为 false，则立即停止迭代。

1
2
3
4
5
6
7
In [68]: from itertools import takewhile

In [69]: list(takewhile(lambda x: x < 5, [1, 3, 6, 2, 1]))
Out[69]: [1, 3]

In [70]: list(takewhile(lambda x: x > 3, [2, 1, 6, 5, 4]))
Out[70]: []
zip_longest

zip_longest 跟zip 类似，但迭代过程会持续到所有可迭代对象的元素都被迭代完。它的形式如下：

1
zip_longest(iter1, iter2, ..., iterN, [fillvalue=None])

如果有指定 fillvalue，则会用其填充缺失的值，否则为 None。

1
2
3
4
5
6
7
8
9
In [75]: from itertools import zip_longest

In [76]: for item in zip_longest('ABCD', 'xy'):
    ...:     print(item)
    ...:
('A', 'x')
('B', 'y')
('C', None)
('D', None)
1
2
3
4
5
6
7
In [77]: for item in zip_longest('ABCD', 'xy', fillvalue='-'):
    ...:     print(item)
    ...:
('A', 'x')
('B', 'y')
('C', '-')
('D', '-')

组合生成器

itertools 模块还提供了多个组合生成器函数，用于求序列的排列、组合等：

  • product
  • permutations
  • combinations
  • combinations_with_replacement
product

product 用于求多个可迭代对象的笛卡尔积，它跟嵌套的 for 循环等价。它的一般使用形式如下：

1
product(iter1, iter2, ... iterN, [repeat=1])

其中，repeat 是一个关键字参数，用于指定重复生成序列的次数

1
2
3
4
5
6
7
8
9
10
11
12
13
In [78]: from itertools import product

In [79]: for item in product('ABCD', 'xy'):
    ...:     print(item)
    ...:
('A', 'x')
('A', 'y')
('B', 'x')
('B', 'y')
('C', 'x')
('C', 'y')
('D', 'x')
('D', 'y')
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
In [80]: list(product('ab', range(3)))
Out[80]: [('a', 0), ('a', 1), ('a', 2), ('b', 0), ('b', 1), ('b', 2)]

In [81]: list(product((0,1), (0,1), (0,1)))
Out[81]:
[(0, 0, 0),
 (0, 0, 1),
 (0, 1, 0),
 (0, 1, 1),
 (1, 0, 0),
 (1, 0, 1),
 (1, 1, 0),
 (1, 1, 1)]

In [82]: list(product('ABC', repeat=2))
Out[82]:
[('A', 'A'),
 ('A', 'B'),
 ('A', 'C'),
 ('B', 'A'),
 ('B', 'B'),
 ('B', 'C'),
 ('C', 'A'),
 ('C', 'B'),
 ('C', 'C')]
permutations

permutations 用于生成一个排列，它的一般使用形式如下：

1
permutations(iterable[, r])

其中，r 指定生成排列的元素的长度，如果不指定，则默认为可迭代对象的元素长度。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
In [83]: from itertools import permutations

In [84]: permutations('ABC', 2)
Out[84]: <itertools.permutations at 0x105e823b8>

In [85]:  list(permutations('ABC', 2))
Out[85]: [('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'C'), ('C', 'A'), ('C', 'B')]

In [86]: list(permutations('ABC'))
Out[86]:
[('A', 'B', 'C'),
 ('A', 'C', 'B'),
 ('B', 'A', 'C'),
 ('B', 'C', 'A'),
 ('C', 'A', 'B'),
 ('C', 'B', 'A')]
combinations

combinations 用于求序列的组合，它的使用形式如下：

1
combinations(iterable, r)

其中，r 指定生成组合的元素的长度。

1
2
3
4
In [87]: from itertools import combinations

In [88]: list(combinations('ABC', 2))
Out[88]: [('A', 'B'), ('A', 'C'), ('B', 'C')]

combinations_with_replacement 和 combinations 类似，但它生成的组合包含自身元素。

1
2
3
4
In [89]: from itertools import combinations_with_replacement

In [90]: list(combinations_with_replacement('ABC', 2))
Out[90]: [('A', 'A'), ('A', 'B'), ('A', 'C'), ('B', 'B'), ('B', 'C'), ('C', 'C')]

itertools 模块提供了很多用于产生多种类型迭代器的函数，它们的返回值不是 list，而是迭代器。

参考文章：高效的 itertools 模块)

知识就是财富
如果您觉得文章对您有帮助, 欢迎请我喝杯水！