【Python】itertools模块，补充：可迭代对象、迭代器

Python中 itertools模块创建高效迭代器、处理序列数据集。

此模块所有函数返回迭代器，可用for循环获取迭代器中的内容，也可用list(...)用列表形式显示内容。

import itertools

[ x for x in dir(itertools) if not x.startswith('_')]
# 结果：
['accumulate', 'chain', 'combinations', 'combinations_with_replacement', 
'compress', 'count', 'cycle', 'dropwhile', 'filterfalse', 'groupby', 
'islice', 'pairwise', 'permutations', 'product', 'repeat', 'starmap',
'takewhile', 'tee', 'zip_longest']

【Python】itertools模块

1、chain 【链接多个可迭代对象】

chain(可迭代对象1, 可迭代对象2, ...)：返回一个迭代器，包含参数中的所有可迭代对象的内容。

from itertools import *

a = [1,2,3,4,5]
b = ["a","b","c","d","e"]

chain(a,b)                          # 结果：
list(chain(a,b))                    # 结果：[1, 2, 3, 4, 5, 'a', 'b', 'c', 'd', 'e']
# 或者
for x in chain(a,b):
    print(x)
# 结果：
1
2
3
4
5
a
b
c
d
e

若不能提前确定所有可迭代对象，或需要懒方法计算，则使用chain.from_iterable(...)。

from itertools import *

# 有yield的函数视为生成器，需用next()方法或for循环才开始执行函数。（控制内存占用，且简洁）
def new_iterable():
    yield [1,2,3]
    yield ["a","b","c"]

for x in chain.from_iterable(new_iterable()):
    print(x)
# 结果：
1
2
3
a
b
c

2、zip_longest 【类似zip】

内置函数zip()。

zip(可迭代对象1, 可迭代对象2, ...)：返回一个迭代器，将多个可迭代对象中的元素，按对应位置一一组成元组形式。若其中一个可迭代对象处理完，则停止。

from itertools import *

a = [1,2,3]
b = ["a","b","c","d","e"]

list(zip(a,b))                  # 结果：[(1, 'a'), (2, 'b'), (3, 'c')]

zip_longest(可迭代对象1, 可迭代对象2, ...)：类似zip。但将所有可迭代对象处理完，若有缺失值则用None填充。可用fillvalue参数指定空缺默认值。

from itertools import *

a = [1,2,3]
b = ["a","b","c","d","e"]

list(zip_longest(a,b))
# 结果：[(1, 'a'), (2, 'b'), (3, 'c'), (None, 'd'), (None, 'e')]
list(zip_longest(a,b,fillvalue=999))
# 结果：[(1, 'a'), (2, 'b'), (3, 'c'), (999, 'd'), (999, 'e')]

3、starmap 【类似map】

内置函数map()。

map(函数，可迭代对象1, 可迭代对象2, ...)：返回一个迭代器，将多个可迭代对象中的元素，依次调用函数处理。若其中一个可迭代对象处理完，则停止。

注：多个可迭代对象，是函数中涉及多个可迭代对象。若函数中只涉及一个可迭代对象，则map参数中只能有一个可迭代对象。

from itertools import *

a = [1,2,3,4,5]
list(map(lambda x:x**3,a))          # 结果：[1, 8, 27, 64, 125]

a = "hello"
b = "world world"
list(map(lambda x,y:(x,y,x+y),a,b))
# 结果：[('h', 'w', 'hw'), ('e', 'o', 'eo'), ('l', 'r', 'lr'), ('l', 'l', 'll'), ('o', 'd', 'od')]

for value in map(lambda x,y:(x,y,x+y),a,b):
    print('{}+{}->{}'.format(*value))
# 或者
# for x,y,z in map(lambda x,y:(x,y,x+y),a,b):
#     print('{}+{}->{}'.format(x,y,z))
# 结果：
h+w->hw
e+o->eo
l+r->lr
l+l->ll
o+d->od

starmap(函数，可迭代对象)：类似map()。区别是：map中的参数是函数和一个或多个可迭代对象，将可迭代对象的每个元素作为单个参数传递给函数即fun(a,b)；而starmap中的参数是函数和一个元组形式的可迭代对象，将可迭代对象的每个元组解包成多个参数传递给函数即fun(*c)。

from itertools import *

c = [('h', 'w'), ('e', 'o'), ('l', 'r'), ('l', 'l'), ('o', 'd')]
list(starmap(lambda x,y:(x,y,x+y),c))
# 结果：[('h', 'w', 'hw'), ('e', 'o', 'eo'), ('l', 'r', 'lr'), ('l', 'l', 'll'), ('o', 'd', 'od')]

for x in starmap(lambda x,y:(x,y,x+y),c):
    print(x)
# 结果：
('h', 'w', 'hw')
('e', 'o', 'eo')
('l', 'r', 'lr')
('l', 'l', 'll')
('o', 'd', 'od')

for value in starmap(lambda x,y:(x,y,x+y),c):
    print('{}+{}->{}'.format(*value))
# 或者
# for x,y,z in starmap(lambda x,y:(x,y,x+y),c):
#     print('{}+{}->{}'.format(x,y,z))
# 结果：
h+w->hw
e+o->eo
l+r->lr
l+l->ll
o+d->od

4、islice 【类似切片】

• islice(可迭代对象, end)：返回一个迭代器，从可迭代对象中获取从0到end（不含）的元素。
• islice(可迭代对象, start, end)：返回一个迭代器，从可迭代对象中获取起始位置start（含）到结束位置end（不含）的元素。若end是None，则一直到可迭代对象结束。
• islice(可迭代对象, start, end, step)：返回一个迭代器，从可迭代对象中获取起始位置start（含）到结束位置end（不含）且步长间隔step的元素。若end是None，则一直到可迭代对象结束。

from itertools import *

list(islice(range(10),5))               # 结果：[0, 1, 2, 3, 4]
list(islice(range(10),5,8))             # 结果：[5, 6, 7]
list(islice(range(10),5,None))          # 结果：[5, 6, 7, 8, 9]
list(islice(range(10),0,8,2))           # 结果：[0, 2, 4, 6]
list(islice(range(10),0,None,2))        # 结果：[0, 2, 4, 6, 8]

 5、accumulate 【累加和】

• accumulate(可迭代对象)：返回一个迭代器，对每个元素依次进行累加和并返回结果。
• accumulate(可迭代对象, initial=num)：返回一个迭代器，在num基础上再对每个元素依次进行累加和并返回结果。
• accumulate(可迭代对象, 函数)：返回一个迭代器，对每个元素依次调用函数并返回结果。

from itertools import *
import operator

list(accumulate(range(10)))      # 结果：[0, 1, 3, 6, 10, 15, 21, 28, 36, 45]
list(accumulate("abcdefg"))      # 结果：['a', 'ab', 'abc', 'abcd', 'abcde', 'abcdef', 'abcdefg']

# 在100的基础上，再进行0-9累加和
list(accumulate(range(10),initial=100))    # 结果：[100, 100, 101, 103, 106, 110, 115, 121, 128, 136, 145]

# 1-5累积相乘
list(accumulate(range(1,6),operator.mul))   # 结果：[1, 2, 6, 24, 120]

6、count 【无限生成连续的整数】

• count()：返回一个迭代器，无限生成连续的整数，从0开始。
• count(start)：返回一个迭代器，无限生成连续的整数，从start开始。
• count(start,step)：返回一个迭代器，无限生成连续的整数，从start开始步长间隔step。

from itertools import *

for x in zip(count(),["ab","cd","ef"]):
    print(x)
# 结果：
(0, 'ab')
(1, 'cd')
(2, 'ef')

list(islice(count(1,4),5))             # 结果：[1, 5, 9, 13, 17]

注解：zip()最短的可迭代对象处理完就停止。count()无限生成连续的整数，但["ab","cd","ef"]长度只有3，处理完就停止。

7、cycle  【无限循环】

cycle(可迭代对象)： 返回一个迭代器，将可迭代对象无限循环。

from itertools import *

list(zip(range(10),cycle("good")))
# 结果：
[(0, 'g'), (1, 'o'), (2, 'o'), (3, 'd'), (4, 'g'), (5, 'o'), (6, 'o'), (7, 'd'), (8, 'g'), (9, 'o')]

list(zip(range(10),cycle(["a","b"])))
# 结果：
[(0, 'a'), (1, 'b'), (2, 'a'), (3, 'b'), (4, 'a'), (5, 'b'), (6, 'a'), (7, 'b'), (8, 'a'), (9, 'b')]

8、repeat 【重复】

 repeat(需重复的数据, num)：返回一个迭代器，重复num次。

repeat(n)：也可为map或zip提供常量值流(a stream of constant values)。

from itertools import *
from operator import pow

list(repeat('abc', 5))         # 结果：['abc', 'abc', 'abc', 'abc', 'abc']
list(repeat(1,5))              # 结果：[1, 1, 1, 1, 1]

# 将列表中的数字求3次幂，num**3
a = [1,2,3]
list(map(pow,a,repeat(3)))     # 结果：[1, 8, 27]

9、product 【笛卡尔积】

• product(可迭代对象1, 可迭代对象2, ...)：返回一个迭代器，将多个可迭代对象中的元素依次一一组成元组。
• product(可迭代对象, repeat=num)：返回一个迭代器，将一个可迭代对象和自身进行笛卡尔积。

from itertools import *

a = [1,2,3]
b = ["a","b"]
c = "hello"
list(product(a,b,c))
# 结果：
[(1, 'a', 'h'), (1, 'a', 'e'), (1, 'a', 'l'), (1, 'a', 'l'), (1, 'a', 'o'),
(1, 'b', 'h'), (1, 'b', 'e'), (1, 'b', 'l'), (1, 'b', 'l'), (1, 'b', 'o'),
(2, 'a', 'h'), (2, 'a', 'e'), (2, 'a', 'l'), (2, 'a', 'l'), (2, 'a', 'o'), 
(2, 'b', 'h'), (2, 'b', 'e'), (2, 'b', 'l'), (2, 'b', 'l'), (2, 'b', 'o'),
(3, 'a', 'h'), (3, 'a', 'e'), (3, 'a', 'l'), (3, 'a', 'l'), (3, 'a', 'o'), 
(3, 'b', 'h'), (3, 'b', 'e'), (3, 'b', 'l'), (3, 'b', 'l'), (3, 'b', 'o')]


a = [1,2,3]
list(product(a, repeat=2))      # 相当于 list(product(a,a))
# 结果：
[(1, 1), (1, 2), (1, 3), (2, 1), (2, 2), (2, 3), (3, 1), (3, 2), (3, 3)]

10、permutations 【指定长度的排列组合】

permutations(可迭代对象)：返回一个迭代器，将元素重新排列且长度为可迭代对象的长度的所有组合，组合以元组形式。

permutations(可迭代对象, n)：返回一个迭代器，将元素重新排列且长度为n的所有组合，组合以元组形式。

from itertools import *

a = [1,2,3]
list(permutations(a))
# 结果：
[(1, 2, 3), (1, 3, 2), (2, 1, 3), (2, 3, 1), (3, 1, 2), (3, 2, 1)]

list(permutations(a,2))
# 结果：
[(1, 2), (1, 3), (2, 1), (2, 3), (3, 1), (3, 2)]

11、combinations 【指定长度的排列组合，没有重复值】

combinations(可迭代对象,r=num)：类似permutations。但组合中没有重复值（重复值为元组中元素相同但元素顺序不同，例如(1,2)和(2,1)为重复值）。

注意：r参数为必需参数，指定长度。

from itertools import *

a = [1,2,3]
list(combinations(a,r=2))      # 结果：[(1, 2), (1, 3), (2, 3)]
list(combinations(a,2))        # 结果：[(1, 2), (1, 3), (2, 3)]

permutations_with_replacement(可迭代对象,r=num)：类似combinations()。但有重复元素的组合（例如(1,1)和(2,2)）。

from itertools import *

a = [1,2,3]
list(combinations_with_replacement(a,r=2))      # 结果：[(1, 1), (1, 2), (1, 3), (2, 2), (2, 3), (3, 3)]

12、pairwise  【依次生成连续的两个元素】

pairwise(可迭代对象)：返回一个迭代器，依次生成连续两个元素的组合。若可迭代对象中的元素小于2个，则返回的迭代器为空。

from itertools import *

a = [1,2,3,4,5,6]
list(pairwise(a))       # 结果：[(1, 2), (2, 3), (3, 4), (4, 5), (5, 6)]

b = "hello"
list(pairwise(b))       # 结果：[('h', 'e'), ('e', 'l'), ('l', 'l'), ('l', 'o')]

c = "w"
list(pairwise(c))       # 结果：[]

13、takewhile【返回满足条件的，舍弃第一个False之后的所有元素】

takewhile(predicate, 可迭代对象)：返回一个迭代器，若predicate返回True，则返回该元素，一旦碰到predicate返回False，则终止。

from itertools import *

a = [3,6,7,8,9]
list(takewhile(lambda x: x%3==0, a))       # 结果：[3, 6]

14、dropwhile【丢弃满足条件的，返回第一个False之后的所有元素】

 dropwhile(predicate, 可迭代对象)：返回一个迭代器，若predicate返回True，则丢弃该元素，一旦碰到predicate返回False，则终止，返回后面的全部元素。

from itertools import *

a = [3,6,7,8,9]
list(dropwhile(lambda x: x%3==0, a))       # 结果：[7, 8, 9]

15、filterfalse 【只返回不满足条件的元素】

内置函数filter(）。

filtler(predicate, 可迭代对象)：返回一个迭代器，只返回满足条件的元素。

from itertools import *

a = [3,6,7,8,9]
list(filter(lambda x: x%3==0, a))          # 结果：[3, 6, 9]

 filterfalse(predicate, 可迭代对象)：返回一个迭代器，只返回不满足条件的元素。

from itertools import *

a = [3,6,7,8,9]
list(filterfalse(lambda x: x%3==0, a))       # 结果：[7, 8]

16、compress  【通过选择器，返回需要的数据】

compress(数据, 选择器)：以选择器作为筛选，来返回数据中符合的数据。

from itertools import *

a = [3,6,7,8,9]
b = cycle([False,True])       # 即 False,True, False,True, False,True...
list(compress(a, b))          # 结果：[6, 8]

17、groupby  【分组】

groupby(可迭代对象, 分组内容)：根据分组内容进行分组，返回一个迭代器，迭代器中的元素是元组，(分组内容, 以列表显示对应的值)。

注意：分组之前，注意排序。

from itertools import *

a = [1,2,2,3,1,2]
b = islice(count(),8)
ab = list(zip(a,b))
ab                               # 结果：[(1, 0), (2, 1), (2, 2), (3, 3), (1, 4), (2, 5)]

for k,v in groupby(ab,lambda x:x[0]):
    print(k,list(v))
# 结果：
1 [(1, 0)]
2 [(2, 1), (2, 2)]
3 [(3, 3)]
1 [(1, 4)]
2 [(2, 5)] 

ab.sort()
ab                               # 结果：[(1, 0), (1, 4), (2, 1), (2, 2), (2, 5), (3, 3)]

for k,v in groupby(ab,lambda x:x[0]):
    print(k,list(v))
# 结果：
1 [(1, 0), (1, 4)]
2 [(2, 1), (2, 2), (2, 5)]
3 [(3, 3)]

18、tee 【一个输入分解成多个相同的输出流】

变量1, 变量2,... = tee(可迭代对象, n=num)：返回元组形式，元组中元素是迭代器。将一个输入分解成num个相同的输出流。默认n=2。类似于UNIX tee工具。

注意：① 使用tee()生成新的迭代器，原来的可迭代对象不应再使用。

          ② tee()生成的迭代器，不是线程安全的，且需大量辅助存储。

from itertools import *

a = [1,2,3,4,5]

m,n = tee(a)                 # 相当于 m,n = tee(a,2) 
list(m)                      # 结果：[0, 1, 2, 3, 4, 5, 6, 7]
list(n)                      # 结果：[0, 1, 2, 3, 4, 5, 6, 7]
list(m) == list(n)           # 结果：True
m is n                       # 结果：False

补充：可迭代对象、迭代器

可迭代对象：

满足__iter__( )方法的对象。即可用for循环遍历的对象。

例如：列表、字符串、文件等。

迭代器：

同时满足__iter__( )方法和__next__( )方法的对象。即可用for循环遍历又可用next()获取下一个值的对象。

可迭代对象不一定是迭代器，迭代器一定是可迭代对象。

可迭代对象转为迭代器：iter(可迭代对象)

判断是否是可迭代对象：isinstance(对象, collections.abc.Iterable)

判断是否是迭代器：isinstance(对象, collections.abc.Iterator)

from collections.abc import Iterator, Iterable

isinstance("hello",Iterable)            # 结果：True
isinstance("hello",Iterator)            # 结果：False

isinstance(iter("hello"),Iterator)      # 结果：True

参考：Pydoc: built-in module itertools