高效的 itertools 模块

(点击 上方蓝字 ，快速关注我们)

来源：FunHacks

funhacks.net/2017/02/13/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：

>>> import itertools

>>>

>>> nums = itertools . count ()

>>> for i in nums :

... if i > 6 :

... break

... print i

...

0

1

2

3

4

5

6

>>> nums = itertools . count ( 10 , 2 ) # 指定开始值和步长

>>> for i in nums :

... if i > 20 :

... break

... print i

...

10

12

14

16

18

20

cycle

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

>>> import itertools

>>>

>>> cycle_strings = itertools . cycle ( 'ABC' )

>>> i = 1

>>> 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：

>>> 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 的使用形式如下：

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

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

>>> from itertools import chain

>>>

>>> for item in chain ([ 1 , 2 , 3 ], [ 'a' , 'b' , 'c' ]) :

... print item

...

1

2

3

a

b

c

chain 还有一个常见的用法：

chain.from_iterable(iterable)

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

>>> from itertools import chain

>>>

>>> string = chain . from_iterable ( 'ABCD' )

>>> string . next ()

'A'

compress

compress 的使用形式如下：

compress(data, selectors)

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

>>> from itertools import compress

>>>

>>> list ( compress ( 'ABCDEF' , [ 1 , 1 , 0 , 1 , 0 , 1 ]))

[ 'A' , 'B' , 'D' , 'F' ]

>>> list ( compress ( 'ABCDEF' , [ 1 , 1 , 0 , 1 ]))

[ 'A' , 'B' , 'D' ]

>>> list ( compress ( 'ABCDEF' , [ True , False , True ]))

[ 'A' , 'C' ]

dropwhile

dropwhile 的使用形式如下：

dropwhile(predicate, iterable)

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

>>> from itertools import dropwhile

>>>

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

[ 6 , 2 , 1 ]

>>>

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

[ 2 , 1 , 6 , 5 , 4 ]

groupby

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

groupby(iterable[, keyfunc])

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

>>> 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' ]

>>>

>>> data = [ 'a' , 'bb' , 'ccc' , 'dd' , 'eee' , 'f' ]

>>> 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' ]

>>>

>>> data = [ 'a' , 'bb' , 'cc' , 'ddd' , 'eee' , 'f' ]

>>> 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 的使用形式如下：

ifilter(function or None, sequence)

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

>>> from itertools import ifilter

>>>

>>> list ( ifilter ( lambda x : x < 6 , range ( 10 )))

[ 0 , 1 , 2 , 3 , 4 , 5 ]

>>>

>>> list ( ifilter ( None , [ 0 , 1 , 2 , 0 , 3 , 4 ]))

[ 1 , 2 , 3 , 4 ]

ifilterfalse

ifilterfalse 的使用形式和 ifilter 类似，它将 iterable 中 function(item) 为 False 的元素组成一个迭代器返回，如果 function 是 None，则返回 iterable 中所有计算为 False 的项。

>>> from itertools import ifilterfalse

>>>

>>> list ( ifilterfalse ( lambda x : x < 6 , range ( 10 )))

[ 6 , 7 , 8 , 9 ]

>>>

>>> list ( ifilter ( None , [ 0 , 1 , 2 , 0 , 3 , 4 ]))

[ 0 , 0 ]

islice

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

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

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

>>> from itertools import count , islice

>>>

>>> list ( islice ([ 10 , 6 , 2 , 8 , 1 , 3 , 9 ], 5 ))

[ 10 , 6 , 2 , 8 , 1 ]

>>>

>>> list ( islice ( count (), 6 ))

[ 0 , 1 , 2 , 3 , 4 , 5 ]

>>>

>>> list ( islice ( count (), 3 , 10 ))

[ 3 , 4 , 5 , 6 , 7 , 8 , 9 ]

>>> list ( islice ( count (), 3 , 10 , 2 ))