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. >>> import itertools
  2. >>>
  3. >>> nums = itertools.count()
  4. >>> for i in nums:
  5. ...     if i > 6:
  6. ...         break
  7. ...     print i
  8. ...
  9. 0
  10. 1
  11. 2
  12. 3
  13. 4
  14. 5
  15. 6
  16. >>> nums = itertools.count(10, 2)    # 指定开始值和步长
  17. >>> for i in nums:
  18. ...     if i > 20:
  19. ...         break
  20. ...     print i
  21. ...
  22. 10
  23. 12
  24. 14
  25. 16
  26. 18
  27. 20

cycle

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

  1. >>> import itertools
  2. >>>
  3. >>> cycle_strings = itertools.cycle('ABC')
  4. >>> i = 1
  5. >>> for string in cycle_strings:
  6. ...     if i == 10:
  7. ...         break
  8. ...     print i, string
  9. ...     i += 1
  10. ...
  11. 1 A
  12. 2 B
  13. 3 C
  14. 4 A
  15. 5 B
  16. 6 C
  17. 7 A
  18. 8 B
  19. 9 C

repeat

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

  1. >>> import itertools
  2. >>>
  3. >>> for item in itertools.repeat('hello world', 3):
  4. ...     print item
  5. ...
  6. hello world
  7. hello world
  8. hello world
  9. >>>
  10. >>> for item in itertools.repeat([1, 2, 3, 4], 3):
  11. ...     print item
  12. ...
  13. [1, 2, 3, 4]
  14. [1, 2, 3, 4]
  15. [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. >>> from itertools import chain
  2. >>>
  3. >>> for item in chain([1, 2, 3], ['a', 'b', 'c']):
  4. ...     print item
  5. ...
  6. 1
  7. 2
  8. 3
  9. a
  10. b
  11. c

chain 还有一个常见的用法：

  1. chain.from_iterable(iterable)

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

  1. >>> from itertools import chain
  2. >>>
  3. >>> string = chain.from_iterable('ABCD')
  4. >>> string.next()
  5. 'A'

compress

compress 的使用形式如下：

  1. compress(data, selectors)

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

  1. >>> from itertools import compress
  2. >>>
  3. >>> list(compress('ABCDEF', [1, 1, 0, 1, 0, 1]))
  4. ['A', 'B', 'D', 'F']
  5. >>> list(compress('ABCDEF', [1, 1, 0, 1]))
  6. ['A', 'B', 'D']
  7. >>> list(compress('ABCDEF', [True, False, True]))
  8. ['A', 'C']

dropwhile

dropwhile 的使用形式如下：

  1. dropwhile(predicate, iterable)

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

  1. >>> from itertools import dropwhile
  2. >>>
  3. >>> list(dropwhile(lambda x: x < 5, [1, 3, 6, 2, 1]))
  4. [6, 2, 1]
  5. >>>
  6. >>> list(dropwhile(lambda x: x > 3, [2, 1, 6, 5, 4]))
  7. [2, 1, 6, 5, 4]

groupby

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

  1. groupby(iterable[, keyfunc])

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

  1. >>> from itertools import groupby
  2. >>>
  3. >>> for key, value_iter in groupby('aaabbbaaccd'):
  4. ...     print key, ':', list(value_iter)
  5. ...
  6. a : ['a', 'a', 'a']
  7. b : ['b', 'b', 'b']
  8. a : ['a', 'a']
  9. c : ['c', 'c']
  10. d : ['d']
  11. >>>
  12. >>> data = ['a', 'bb', 'ccc', 'dd', 'eee', 'f']
  13. >>> for key, value_iter in groupby(data, len):    # 使用 len 函数作为分组函数
  14. ...     print key, ':', list(value_iter)
  15. ...
  16. 1 : ['a']
  17. 2 : ['bb']
  18. 3 : ['ccc']
  19. 2 : ['dd']
  20. 3 : ['eee']
  21. 1 : ['f']
  22. >>>
  23. >>> data = ['a', 'bb', 'cc', 'ddd', 'eee', 'f']
  24. >>> for key, value_iter in groupby(data, len):
  25. ...     print key, ':', list(value_iter)
  26. ...
  27. 1 : ['a']
  28. 2 : ['bb', 'cc']
  29. 3 : ['ddd', 'eee']
  30. 1 : ['f']

ifilter

ifilter 的使用形式如下：

  1. ifilter(function or None, sequence)

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

  1. >>> from itertools import ifilter
  2. >>>
  3. >>> list(ifilter(lambda x: x < 6, range(10)))
  4. [0, 1, 2, 3, 4, 5]
  5. >>>
  6. >>> list(ifilter(None, [0, 1, 2, 0, 3, 4]))
  7. [1, 2, 3, 4]

ifilterfalse

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

  1. >>> from itertools import ifilterfalse
  2. >>>
  3. >>> list(ifilterfalse(lambda x: x < 6, range(10)))
  4. [6, 7, 8, 9]
  5. >>>
  6. >>> list(ifilter(None, [0, 1, 2, 0, 3, 4]))
  7. [0, 0]

islice

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

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

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

  1. >>> from itertools import count, islice
  2. >>>
  3. >>> list(islice([10, 6, 2, 8, 1, 3, 9], 5))
  4. [10, 6, 2, 8, 1]
  5. >>>
  6. >>> list(islice(count(), 6))
  7. [0, 1, 2, 3, 4, 5]
  8. >>>
  9. >>> list(islice(count(), 3, 10))
  10. [3, 4, 5, 6, 7, 8, 9]
  11. >>> list(islice(count(), 3, 10 ,2))
  12. [3, 5, 7, 9]

imap

imap 类似 map 操作，它的使用形式如下：

  1. imap(func, iter1, iter2, iter3, ...)

imap 返回一个迭代器，元素为 func(i1, i2, i3, ...)，i1，i2 等分别来源于 iter, iter2。

  1. >>> from itertools import imap
  2. >>>
  3. >>> imap(str, [1, 2, 3, 4])
  4. <itertools.imap object at 0x10556d050>
  5. >>>
  6. >>> list(imap(str, [1, 2, 3, 4]))
  7. ['1', '2', '3', '4']
  8. >>>
  9. >>> list(imap(pow, [2, 3, 10], [4, 2, 3]))
  10. [16, 9, 1000]

tee

tee 的使用形式如下：

  1. tee(iterable [,n])

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

  1. >>> from itertools import tee
  2. >>>
  3. >>> tee('abcd')   # n 默认为 2，创建两个独立的迭代器
  4. (<itertools.tee object at 0x1049957e8>, <itertools.tee object at 0x104995878>)
  5. >>>
  6. >>> iter1, iter2 = tee('abcde')
  7. >>> list(iter1)
  8. ['a', 'b', 'c', 'd', 'e']
  9. >>> list(iter2)
  10. ['a', 'b', 'c', 'd', 'e']
  11. >>>
  12. >>> tee('abc', 3)  # 创建三个独立的迭代器
  13. (<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. >>> from itertools import takewhile
  2. >>>
  3. >>> list(takewhile(lambda x: x < 5, [1, 3, 6, 2, 1]))
  4. [1, 3]
  5. >>> list(takewhile(lambda x: x > 3, [2, 1, 6, 5, 4]))
  6. []

izip

izip 用于将多个可迭代对象对应位置的元素作为一个元组，将所有元组『组成』一个迭代器，并返回。它的使用形式如下：

  1. izip(iter1, iter2, ..., iterN)

如果某个可迭代对象不再生成值，则迭代停止。

  1. >>> from itertools import izip
  2. >>> 
  3. >>> for item in izip('ABCD', 'xy'):
  4. ...     print item
  5. ...
  6. ('A', 'x')
  7. ('B', 'y')
  8. >>> for item in izip([1, 2, 3], ['a', 'b', 'c', 'd', 'e']):
  9. ...     print item
  10. ...
  11. (1, 'a')
  12. (2, 'b')
  13. (3, 'c')

izip_longest

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

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

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

  1. >>> from itertools import izip_longest
  2. >>>
  3. >>> for item in izip_longest('ABCD', 'xy'):
  4. ...     print item
  5. ...
  6. ('A', 'x')
  7. ('B', 'y')
  8. ('C', None)
  9. ('D', None)
  10. >>>
  11. >>> for item in izip_longest('ABCD', 'xy', fillvalue='-'):
  12. ...     print item
  13. ...
  14. ('A', 'x')
  15. ('B', 'y')
  16. ('C', '-')
  17. ('D', '-')

组合生成器

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

  • product
  • permutations
  • combinations
  • combinations_with_replacement

product

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

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

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

  1. >>> from itertools import product
  2. >>>
  3. >>> for item in product('ABCD', 'xy'):
  4. ...     print item
  5. ...
  6. ('A', 'x')
  7. ('A', 'y')
  8. ('B', 'x')
  9. ('B', 'y')
  10. ('C', 'x')
  11. ('C', 'y')
  12. ('D', 'x')
  13. ('D', 'y')
  14. >>>
  15. >>> list(product('ab', range(3)))
  16. [('a', 0), ('a', 1), ('a', 2), ('b', 0), ('b', 1), ('b', 2)]
  17. >>>
  18. >>> list(product((0,1), (0,1), (0,1)))
  19. [(0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1), (1, 0, 0), (1, 0, 1), (1, 1, 0), (1, 1, 1)]
  20. >>>
  21. >>> list(product('ABC', repeat=2))
  22. [('A', 'A'), ('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'B'), ('B', 'C'), ('C', 'A'), ('C', 'B'), ('C', 'C')]
  23. >>>

permutations

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

  1. permutations(iterable[, r])

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

  1. >>> from itertools import permutations
  2. >>>
  3. >>> permutations('ABC', 2)
  4. <itertools.permutations object at 0x1074d9c50>
  5. >>>
  6. >>> list(permutations('ABC', 2))
  7. [('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'C'), ('C', 'A'), ('C', 'B')]
  8. >>>
  9. >>> list(permutations('ABC'))
  10. [('A', 'B', 'C'), ('A', 'C', 'B'), ('B', 'A', 'C'), ('B', 'C', 'A'), ('C', 'A', 'B'), ('C', 'B', 'A')]
  11. >>>

combinations

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

  1. combinations(iterable, r)

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

  1. >>> from itertools import combinations
  2. >>>
  3. >>> list(combinations('ABC', 2))
  4. [('A', 'B'), ('A', 'C'), ('B', 'C')]

combinations_with_replacement

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

  1. >>> from itertools import combinations_with_replacement
  2. >>>
  3. >>> list(combinations_with_replacement('ABC', 2))
  4. [('A', 'A'), ('A', 'B'), ('A', 'C'), ('B', 'B'), ('B', 'C'), ('C', 'C')]

小结

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

参考链接