高效迭代器 itertools

Python 的标准函数“itertools”是一个针对可迭代对象进行处理的函数，由于是 Python 内建的标准函数，因此处理数据的速度，比自己撰写程序来迭代每个项目还来得迅速，这篇文章将会介绍 itertools 的常用方法。

本篇使用的 Python 版本为 3.7.12，所有范例可使用 Google Colab 实作，不用安装任何软件 ( 参考：使用 Google Colab )

itertools 常用方法

itertools 包含一系列用来产生不同类型迭代器的方法，这些方法都会回传一个迭代器，可以透过循环的方式取值，也可以使用 next() 来取值，itertools 有以下几种类型与常用的方法 ( 参考 Python 官方文件：time )：

无限迭代器

有限迭代器

排列组合迭代器

import itertools

要使用 time 必须先 import itertools 模组，或使用 from 的方式，单独 import 特定的类型。

import itertools
from itertools import count

count(start, step)

itertools.count(start, step) 执行后会回传从 start 开始，以 step 间隔的无限项目迭代器，step 默认 1，不填入 step 会自动使用 1，由于 count 回传的迭代器是一个无限循环的产生器，所以如果使用循环，则必须要自己撰写判断式停止。

import itertools
a = itertools.count(1)         # 設定 a 從 1 開始，間隔 1 無限循環
b = itertools.count(5, 2)      # 設定 b 從 5 開始，間隔 2 無限循環
for i in a:
    print(i, end=' ')          # 1 2 3 4 5 6 7 印出 a 裡的每個項目
    if i>6: break              # 如果超過 6 就停止
print()
for i in b:
    print(i, end=' ')          # 5 7 9 11 13 15 17 19 21 印出 b 裡的每個項目
    if i>20: break             # 如果超過 20 就停止

因为是产生器，所以能够使用 next() 读取，下方程序码就会依序读取每个项目。

import itertools
a = itertools.count(1)
print(next(a))   # 1
print(next(a))   # 2
print(next(a))   # 3
print(next(a))   # 4
print(next(a))   # 5

cycle(iterable)

itertools.cycle(iterable) 执行后会回传一个迭代器，内容是迭代对象中所有元素的无限重复，下方的程序码执行后，使用 for 循环印出 20 个项目，这 20 个项目是使用 ABC 三个字母无限循环所产生。

import itertools
a = itertools.cycle('ABC')
for i in range(20):
    print(next(a), end=' ')    # A B C A B C A B C A B C A B C A B C A B

repeat(object, times)

itertools.repeat(object, times) 执行后会回传一个迭代器，内容是迭代对象的无限重复，times 默认 None 为无限重复，可指定 times 设定重复次数，下方的程序码执行后，会印出五次 ABC。

import itertools
a = itertools.repeat('ABC', 5)
for i in a:
    print(i, end=' ')   # ABC ABC ABC ABC ABC

accumulate(p, fn)

itertools.accumulate(p, fn) 如果不设定 fn，执行后默认回传累积加总值，如果设定 fn，则会回传运算函数的累积结果值，下方的例子，b 的每个项目会是 a 的每个项目加上前面每个项目的数值，c 的每个项目则是 a 的每个项目乘以前面每个项目的数值。

import itertools
a = [1,2,3,4,5,6,7,8,9]
b = itertools.accumulate(a)
c = itertools.accumulate(a, lambda x,y: x*y)  # 使用 lambda 匿名函式
for i in b:
    print(i, end=' ')   # 1 3 6 10 15 21 28 36 45
print()
for i in c:
    print(i, end=' ')   # 1 2 6 24 120 720 5040 40320 362880

chain(p, q...)

itertools.chain(p, q...) 可以将多个可迭代对象，串连为单一个可迭代对象，下方的程序码执行后，会将 a、b、c 三个可迭代对象，组合成单一个可迭代的对象。

import itertools
a = 'abc'
b = 'xyz'
c = [1,2,3]
d = itertools.chain(a, b, c)
for i in d:
    print(i, end=' ')   # a b c x y z 1 2 3

compress(data, selectors)

itertools.compress(data, selectors) 会将 data 和 selectors 比对，当 selectors 的某个元素为 true 时，则保留 data 对应位置的元素，下方的程序码执行后，会以 b 为 True 的项目位置 1、2、5 挑出 a 的项目，就会回传 A、B、E。

import itertools
a = 'ABCDEFG'
b = [1,1,0,0,1]
c = itertools.compress(a, b)
for i in c:
    print(i, end=' ')   # A B E

dropwhile(func, iter)

itertools.dropwhile(pred, seq) 里的 func 是函数，iter 是可迭代对象。透过 func 函数计算可迭代的项目，如果结果为 True 就舍弃该项目，如果是 False 就回传该项目以及后方所有项目，下方的程序码执行后，b 的结果因为在 5 的位置为 Fasle，所以会回传 5 和后方所有项目，而 c 的结果因为在 1 的位置就 False，所以就会回传 1 和后方所有项目。

import itertools
a = [1,2,3,4,5,6,7,8,9]
b = itertools.dropwhile(lambda x: x<5, a)
c = itertools.dropwhile(lambda x: x>5, a)
for i in b:
    print(i, end=' ')    # 5 6 7 8 9
print()
for i in c:
    print(i, end=' ')    # 1 2 3 4 5 6 7 8 9

filterfalse(func, iter)

itertools.filterfalse(pred, seq) 里的 func 是函数，iter 是可迭代对象。透过 func 函数计算可迭代的项目，如果结果为 False 就回传该项目，下方的程序码执行后，b 会回传大于等于 5 的数值，c 会回传小于等于 5 的数值。

import itertools
a = [1,2,3,4,5,6,7,8,9]
b = itertools.filterfalse(lambda x: x<5, a)
c = itertools.filterfalse(lambda x: x>5, a)
for i in b:
    print(i, end=' ')    # 5 6 7 8 9
print()
for i in c:
    print(i, end=' ')    # 1 2 3 4 5

takewhile(func, iter)

itertools.takewhile(func, iter) 里的 func 是函数，iter 是可迭代对象。透过 func 函数计算可迭代的项目，如果结果为 True 就回传该项目，如果遇到 False 就会停止，下方的程序码执行后，b 会回传小于 5 的数值，c 则不会回传数据，因为一开始遇到 False 就会停止。

import itertools
a = [1,2,3,4,5,6,7,8,9]
b = itertools.takewhile(lambda x: x<5, a)
c = itertools.takewhile(lambda x: x>5, a)
for i in b:
    print(i, end=' ')    # 1 2 3 4
print()
for i in c:
    print(i, end=' ')    # 沒有結果，因為一開始就遇到 False

groupby(iter, key)

itertools.groupby(iter, key) 会将相邻的项目，按照 key 的设定集合成一组，如果没有 key 就会把“相邻且相同”的项目分成同一组，默认 key 为 None，下方的程序码执行后，b 会将“相邻且相同”的项目分成同一组 ( 因为没有设定key )，而 c 会把转换成大写之后“相邻且相同”的项目分成同一组。

import itertools
a = 'AAaBbbCcC'
b = itertools.groupby(a)
for key, val in b:
    print(key, list(val))

# A ['A', 'A']
# a ['a']
# B ['B']
# b ['b', 'b']
# C ['C']
# c ['c']
# C ['C']

c = itertools.groupby(a, lambda x: x.upper())   # 轉換成大寫後分組
for key, val in c:
    print(key, list(val))

# A ['A', 'A', 'a']
# B ['B', 'b', 'b']
# C ['C', 'c', 'C']

下面的程序会先把一个串列，按照项目的“长度”排序，排序后会按照 key 的设定，把同样长度的分成同一组。

import itertools
d = ['dd','ddd','aa','bbbbb','a','ccc','ee']
dd = sorted(d, key=len)          # 按照長度進行排序
e = itertools.groupby(dd, len)   # 按照長度進行分組
for key, val in e:
    print(key, list(val))

# 1 ['a']
# 2 ['dd', 'aa', 'ee']
# 3 ['ddd', 'ccc']
# 5 ['bbbbb']

islice(iter, start, stop, step)

itertools.islice(iter, start, stop, step) 会针对迭代的对象进行切片动作，start 默认 0，表示开始的顺序，stop 默认 None 表示最后一项，step 表示取值的间隔，下面的项目执行后，会从第二个项目开始，取出间隔 2 的项目。

import itertools
a = '123456789'
b = itertools.islice(a, 2, len(a), 2)
for i in b:
    print(i, end=' ')   # 3 5 7 9

starmap(func, iter)

itertools.starmap(func, iter) 会将可迭代的对象透过 func 函数运算后，产生新的项目，类似 map() 的用法，下方的程序执行后，会把 a 所有的字母变成两倍。

import itertools
a = 'abcdefg'
b = itertools.starmap(lambda x: x*2, a)
for i in b:
    print(i, end=' ')

tee(iter, n)

itertools.tee(iter, n) 会将原本的迭代对象，产生 n 个独立的迭代器，下方的程序执行后，会将 a 变成三个独立的迭代器。

import itertools
a = 'abcde'
b = itertools.tee(a, 3)
for i in b:
    print(list(i))

# ['a', 'b', 'c', 'd', 'e']
# ['a', 'b', 'c', 'd', 'e']
# ['a', 'b', 'c', 'd', 'e']

zip_longest(*iters, fillvalue)

itertools.zip_longest(iters, fillvalue) 会将多个可迭代对象合并，如果遇到不足的项目，默认补上 None，也可设定 fillvalue 补值，下方的程序码执行后，会将 a 和 b 组合成可迭代对象，不足的项目以问号 ? 补充。

import itertools
a = 'abcde'
b = '123'
c = itertools.zip_longest(a, b, fillvalue='?')
for i in c:
    print(i)

product(p, q...)

itertools.product(p, q...) 会将多个可迭代对象，回传笛卡尔积的结果 ( 参考笛卡儿积 )，下方程序执行后，会回传 a 和 b 的笛卡尔积。

import itertools
a = 'abc'
b = '123'
r = itertools.product(a, b)
for i in r:
    print(*i)

# a 1
# a 2
# a 3
# b 1
# b 2
# b 3
# c 1
# c 2
# c 3

permutations(p)

itertools.permutations(p) 会将可迭代对象按照顺序，进行所有不重复的排列组合，并回传最后的结果，下方的程序码执行后，会回传 abc 三个字母所有的排列组合。

import itertools
a = 'abc'
r = itertools.permutations(a)
for i in r:
    print(*i)

# a b c
# a c b
# b a c
# b c a
# c a b
# c b a

combinations(p, r)

itertools.combinations(p, r) 会将可迭代对象按照顺序，进行指定长度的组合，并回传最后的结果，和 permutations 不同的是，permutations 是所有的排列组合 ( 项目的前后顺序不同视为不同项目 )，combinations 则是组合，如果内容元素相同，就算排列顺序不同，仍视为相同项目。

下方的程序码执行后，会以 3 的长度印出对应的组合。

import itertools
a = 'abcde'
r = itertools.combinations(a,3)
for i in r:
    print(*i)

# a b c
# a b d
# a b e
# a c d
# a c e
# a d e
# b c d
# b c e
# b d e
# c d e

combinations_with_replacement(p, r)

itertools.combinations_with_replacement(p, r) 和 combinations 相同，只是组合里会包含自己的项目，下方程序码执行后，会以 2 的长度印出对应的组合。

import itertools
a = 'abcde'
r = itertools.combinations_with_replacement(a,2)
for i in r:
    print(*i)
'''
a a
a b
a c
a d
a e
b b
b c
b d
b e
c c
c d
c e
d d
d e
e e
'''

意见反馈