高效的 itertools 模块
我们知道,迭代器的特点是:惰性求值(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
17
18
19
20
21
22
23
24
25
26
27
|
import itertools
>>>
for i
in nums:
if i >
6:
break
print i
...
0
1
2
3
4
5
6
10,
2)
# 指定开始值和步长
nums = itertools.count(
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
|
import itertools
>>>
'ABC')
cycle_strings = itertools.cycle(
1
i =
for string
in cycle_strings:
if i ==
10:
break
print i, string
1
i +=
...
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
7
8
9
10
11
12
13
14
15
|
import itertools
>>>
for item
in itertools.repeat(
'hello world',
3):
print item
...
hello world
hello world
hello world
>>>
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()
- ifilter()
- ifilterfalse()
- islice()
- imap()
- starmap()
- tee()
- takewhile()
- izip()
- izip_longest()
chain
chain
的使用形式如下:
1
|
chain(iterable1, iterable2, iterable3, ...)
|
chain
接收多个可迭代对象作为参数,将它们『连接』起来,作为一个新的迭代器返回。
1
2
3
4
5
6
7
8
9
10
11
|
from itertools
import chain
>>>
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
|
from itertools
import chain
>>>
'ABCD')
string = chain.from_iterable(
'A'
|
compress
compress
的使用形式如下:
1
|
compress(data, selectors)
|
compress
可用于对数据进行筛选,当 selectors 的某个元素为 true 时,则保留 data 对应位置的元素,否则去除:
1
2
3
4
5
6
7
8
|
from itertools
import compress
>>>
'ABCDEF', [
1,
1,
0,
1,
0,
1]))
list(compress(
[
'A',
'B',
'D',
'F']
'ABCDEF', [
1,
1,
0,
1]))
list(compress(
[
'A',
'B',
'D']
'ABCDEF', [
True,
False,
True]))
list(compress(
[
'A',
'C']
|
dropwhile
dropwhile
的使用形式如下:
1
|
dropwhile(predicate, iterable)
|
其中,predicate 是函数,iterable 是可迭代对象。对于 iterable 中的元素,如果 predicate(item) 为 true,则丢弃该元素,否则返回该项及所有后续项。
1
2
3
4
5
6
7
|
from itertools
import dropwhile
>>>
lambda x: x <
5, [
1,
3,
6,
2,
1]))
list(dropwhile(
[
6,
2,
1]
>>>
lambda x: x >
3, [
2,
1,
6,
5,
4]))
list(dropwhile(
[
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
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
|
from itertools
import groupby
>>>
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']
>>>
'a',
'bb',
'ccc',
'dd',
'eee',
'f']
data = [
for key, value_iter
in groupby(data, len):
# 使用 len 函数作为分组函数
print key,
':', list(value_iter)
...
1 : [
'a']
2 : [
'bb']
3 : [
'ccc']
2 : [
'dd']
3 : [
'eee']
1 : [
'f']
>>>
'a',
'bb',
'cc',
'ddd',
'eee',
'f']
data = [
for key, value_iter
in groupby(data, len):
print key,
':', list(value_iter)
...
1 : [
'a']
2 : [
'bb',
'cc']
3 : [
'ddd',
'eee']
1 : [
'f']
|
ifilter
ifilter
的使用形式如下:
1
|
ifilter(function or None, sequence)
|
将 iterable 中 function(item) 为 True 的元素组成一个迭代器返回,如果 function 是 None,则返回 iterable 中所有计算为 True 的项。
1
2
3
4
5
6
7
|
from itertools
import ifilter
>>>
lambda x: x <
6, range(
10)))
list(ifilter(
[
0,
1,
2,
3,
4,
5]
>>>
None, [
0,
1,
2,
0,
3,
4]))
list(ifilter(
[
1,
2,
3,
4]
|
ifilterfalse
的使用形式和 ifilter
类似,它将 iterable 中 function(item) 为 False 的元素组成一个迭代器返回,如果 function 是 None,则返回 iterable 中所有计算为 False 的项。
1
2
3
4
5
6
7
|
from itertools
import ifilterfalse
>>>
lambda x: x <
6, range(
10)))
list(ifilterfalse(
[
6,
7,
8,
9]
>>>
None, [
0,
1,
2,
0,
3,
4]))
list(ifilter(
[
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
|
from itertools
import count, islice
>>>
10,
6,
2,
8,
1,
3,
9],
5))
list(islice([
[
10,
6,
2,
8,
1]
>>>
6))
list(islice(count(),
[
0,
1,
2,
3,
4,
5]
>>>
3,
10))
list(islice(count(),
[
3,
4,
5,
6,
7,
8,
9]
3,
10 ,
2))
list(islice(count(),
[
3,
5,
7,
9]
|
imap
imap
类似 map
操作,它的使用形式如下:
1
|
imap(func, iter1, iter2, iter3, ...)
|
imap
返回一个迭代器,元素为 func(i1, i2, i3, ...)
,i1
,i2
等分别来源于 iter
, iter2
。
1
2
3
4
5
6
7
8
9
10
|
from itertools
import imap
>>>
1,
2,
3,
4])
imap(str, [
<itertools.imap object at
0x10556d050>
>>>
1,
2,
3,
4]))
list(imap(str, [
[
'1',
'2',
'3',
'4']
>>>
2,
3,
10], [
4,
2,
3]))
list(imap(pow, [
[
16,
9,
1000]
|
tee
tee
的使用形式如下:
1
|
tee(iterable [,n])
|
tee
用于从 iterable 创建 n 个独立的迭代器,以元组的形式返回,n 的默认值是 2。
1
2
3
4
5
6
7
8
9
10
11
12
13
|
from itertools
import tee
>>>
'abcd')
# n 默认为 2,创建两个独立的迭代器
tee(
(<itertools.tee object at
0x1049957e8>, <itertools.tee object at
0x104995878>)
>>>
'abcde')
iter1, iter2 = tee(
[
'a',
'b',
'c',
'd',
'e']
[
'a',
'b',
'c',
'd',
'e']
>>>
'abc',
3)
# 创建三个独立的迭代器
tee(
(<itertools.tee object at
0x104995998>, <itertools.tee object at
0x1049959e0>, <itertools.tee object at
0x104995a28>)
|
takewhile
takewhile
的使用形式如下:
1
|
takewhile(predicate, iterable)
|
其中,predicate 是函数,iterable 是可迭代对象。对于 iterable 中的元素,如果 predicate(item) 为 true,则保留该元素,只要 predicate(item) 为 false,则立即停止迭代。
1
2
3
4
5
6
|
from itertools
import takewhile
>>>
lambda x: x <
5, [
1,
3,
6,
2,
1]))
list(takewhile(
[
1,
3]
lambda x: x >
3, [
2,
1,
6,
5,
4]))
list(takewhile(
[]
|
izip
izip
用于将多个可迭代对象对应位置的元素作为一个元组,将所有元组『组成』一个迭代器,并返回。它的使用形式如下:
1
|
izip(iter1, iter2, ..., iterN)
|
如果某个可迭代对象不再生成值,则迭代停止。
1
2
3
4
5
6
7
8
9
10
11
12
13
|
from itertools
import izip
for item
in izip(
'ABCD',
'xy'):
print item
...
(
'A',
'x')
(
'B',
'y')
for item
in izip([
1,
2,
3], [
'a',
'b',
'c',
'd',
'e']):
print item
...
(
1,
'a')
(
2,
'b')
(
3,
'c')
|
izip_longest
跟 izip
类似,但迭代过程会持续到所有可迭代对象的元素都被迭代完。它的形式如下:
1
|
izip_longest(iter1, iter2, ..., iterN, [fillvalue=None])
|
如果有指定 fillvalue,则会用其填充缺失的值,否则为 None。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
|
from itertools
import izip_longest
>>>
for item
in izip_longest(
'ABCD',
'xy'):
print item
...
(
'A',
'x')
(
'B',
'y')
(
'C',
None)
(
'D',
None)
>>>
for item
in izip_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
14
15
16
17
18
19
20
21
22
23
|
from itertools
import product
>>>
for item
in product(
'ABCD',
'xy'):
print item
...
(
'A',
'x')
(
'A',
'y')
(
'B',
'x')
(
'B',
'y')
(
'C',
'x')
(
'C',
'y')
(
'D',
'x')
(
'D',
'y')
>>>
'ab', range(
3)))
list(product(
[(
'a',
0), (
'a',
1), (
'a',
2), (
'b',
0), (
'b',
1), (
'b',
2)]
>>>
0,
1), (
0,
1), (
0,
1)))
list(product((
[(
0,
0,
0), (
0,
0,
1), (
0,
1,
0), (
0,
1,
1), (
1,
0,
0), (
1,
0,
1), (
1,
1,
0), (
1,
1,
1)]
>>>
'ABC', repeat=
2))
list(product(
[(
'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
|
from itertools
import permutations
>>>
'ABC',
2)
permutations(
<itertools.permutations object at
0x1074d9c50>
>>>
'ABC',
2))
list(permutations(
[(
'A',
'B'), (
'A',
'C'), (
'B',
'A'), (
'B',
'C'), (
'C',
'A'), (
'C',
'B')]
>>>
'ABC'))
list(permutations(
[(
'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
|
from itertools
import combinations
>>>
'ABC',
2))
list(combinations(
[(
'A',
'B'), (
'A',
'C'), (
'B',
'C')]
|
combinations_with_replacement
和 combinations
类似,但它生成的组合包含自身元素。
1
2
3
4
|
from itertools
import combinations_with_replacement
>>>
'ABC',
2))
list(combinations_with_replacement(
[(
'A',
'A'), (
'A',
'B'), (
'A',
'C'), (
'B',
'B'), (
'B',
'C'), (
'C',
'C')]
|
小结
- itertools 模块提供了很多用于产生多种类型迭代器的函数,它们的返回值不是 list,而是迭代器。