Python中itertools 模組的用法

在 Python 中，迭代器是一種非常好用的資料結構，其最大的優勢就是延遲生成，按需使用，從而大大提高程式的執行效率。而 itertools 作為 Python 的內建模組，就為我們提供了一套非常有用的用於操作可迭代物件的函式。

常用功能

1.count 功能詳解

count(start=0,step=1) 函式有兩個引數，其中 step 是預設引數，可選的，預設值為 1。 該函式返回一個新的迭代器，從 start 開始，返回以 step 為步長的均勻間隔的值。

import itertools
x = itertools.count(1,2)
for k in x:
	print(k, end=", ")

# 輸出結果如下 無窮無盡
1, 3, 5, 7, 9, 11, 13, 15, ...

2.cycle 功能詳解

cycle(iterable) 該函式會把接收到的序列無限重複下去。

import itertools
x = itertools.cycle("XYZ")
for k in x:
	print(k, end = ", ")
  
# 輸出結果如下 無窮無盡
X, Y, Z, X, Y, Z, X, Y, Z, ...

注意，該函式可能需要相當大的輔助空間（取決於 iterable 的長度）。

3.repeat 功能詳解

repeat(object, times) 該函式建立一個迭代器，不斷的重複 object，當然如果指定 times 的話，則只會重複 times 次。

import itertools
x = itertools.repeat("XYZ")
for k in x:
	print(k, end = ", ")
  
# 輸出結果如下 無窮無盡
XYZ, XYZ, XYZ, XYZ, XYZ, XYZ, ...

import itertools
x = itertools.repeat("XYZ", 3)
print(list(x))
#學習中遇到問題沒人解答？小編建立了一個Python學習交流群：153708845
# 輸出結果如下 只會輸出三次
['XYZ', 'XYZ', 'XYZ']

注意：無限迴圈迭代器只有在 for 迴圈中才會不斷的生成元素，如果只是建立一個迭代器物件，則不會事先生成無限個元素。

4.chain 功能詳解

chain(*iterables) 該函式建立一個新的迭代器，會將引數中的所有迭代器全包含進去。

import itertools
x = itertools.chain("abc", "xyz")
print(list(x))

# 輸出結果如下
['a', 'b', 'c', 'x', 'y', 'z']

5.groupby 功能詳解

groupby(iterable, key=None) 分組函式，將 key 函式作用於序列的各個元素。根據 key 函式的返回值將擁有相同返回值的元素分到一個新的迭代器。類似於 SQL 中的 GROUP BY 操作，唯一不同的是該函式對序列的順序有要求，因為當 key 函式的返回值改變時，迭代器就會生成一個新的分組。因此在使用該函式之前需要先使用同一個排序函式對該序列進行排序操作。

import itertools
def sortBy(score):
	if score > 80:
		return "A"
	elif score >= 60:
		return "B"
	else:
		return "C"

scores = [81, 82, 84, 76, 64, 78, 59, 44, 55, 89]
for m, n in itertools.groupby(scores, key=sortBy):
	print(m, list(n))

# 輸出結果如下
A [81, 82, 84]
B [76, 64, 78]
C [59, 44, 55]
A [89]

我們可以看到，該函式根據我們自定義的排序函式 sortBy 將列表中的元素進行了分組操作，只是我們發現最後一個怎麼多了一個 A 的分組呢，這就是我們上面說所得「當 key 函式的返回值改變時，迭代器就會生成一個新的分組」。所以，我們需要事先對列表用 sortBy 函式排一下序。

scores = [81, 82, 84, 76, 64, 78, 59, 44, 55, 89]
scores = sorted(scores, key=sortBy)
for m, n in itertools.groupby(scores, key=sortBy):
	print(m, list(n))

# 輸出結果如下
A [81, 82, 84]
B [76, 64, 78]
C [59, 44, 55]
A [89]

6.compress 功能詳解

compress(data, selectors) 該函式功能很簡單，就是根據 selectors 中的值判斷是否保留 data 中對應位置的值。

import itertools
data = [81, 82, 84, 76, 64, 78]
tf = [1,1,0,1,1,0]
print(list(itertools.compress(data, tf)))

# 輸出結果如下
[81, 82, 76, 64]

7.dropwhile 功能詳解

dropwhile(predicate, iterable) 建立一個迭代器，從 predicate 首次為 false 時開始迭代元素。

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

# 輸出結果如下
[5, 7, 4, 2, 1]

由以上得知，即使 predicate 首次為 false 後面的元素不滿足 predicate 也同樣會被迭代。

8.filterfalse 功能詳解

filterfalse(predicate, iterable) 建立一個迭代器，返回 iterable 中 predicate 為 false 的元素。

import itertools
x = itertools.filterfalse(lambda x: x < 5, [1,3,5,7,4,2,1])
print(list(x))

# 輸出結果如下
[5, 7]

9.islice 功能詳解

islice(iterable, start, stop[, step]) 對 iterable 進行切片操作。從 start 開始到 stop 截止，同時支援以步長為 step 的跳躍。

import itertools
print(list(itertools.islice('123456789', 2)))
print(list(itertools.islice('123456789', 2, 4)))
print(list(itertools.islice('123456789', 2, None)))
print(list(itertools.islice('123456789', 0, None, 2)))

# 輸出結果如下
['1', '2']
['3', '4']
['3', '4', '5', '6', '7', '8', '9']
['1', '3', '5', '7', '9']

10.starmap 功能詳解

starmap(function, iterable) 從可迭代物件中獲取引數來執行該函式。

import itertools
print(list(itertools.starmap(pow,[(2,10), (3,3)])))

# 輸出結果如下
[1024, 27]

11.takewhile 功能詳解

takewhile(predicate, iterable) 建立一個迭代器，遇到 predicate 為 false 則停止迭代元素。與 dropwhile 完全相反。

import itertools
x = itertools.takewhile(lambda x: x < 5, [1,3,5,7,4,2,1])
print(list(x))

# 輸出結果如下
[1, 3]

12.product 功能詳解

product(*iterables, repeat=1) 輸出可迭代物件的笛卡爾積，有點類似於巢狀迴圈。其中 repeat 可以設定迴圈次數。

import itertools
print(list(itertools.product("ab", "12")))
print(list(itertools.product("ab", "ab")))
print(list(itertools.product("ab", repeat=2)))

# 輸出結果如下
[('a', '1'), ('a', '2'), ('b', '1'), ('b', '2')]
[('a', 'a'), ('a', 'b'), ('b', 'a'), ('b', 'b')]
[('a', 'a'), ('a', 'b'), ('b', 'a'), ('b', 'b')]

13.permutations 功能詳解

permutations(iterable, r=None) 返回 iterable 中長度為 r 的所有排列。預設值 r 為 iterable 的長度。即使元素的值相同，不同位置的元素也被認為是不同的。

import itertools
print(list(itertools.permutations("aba", r=2)))

# 輸出結果如下
[('a', 'b'), ('a', 'a'), ('b', 'a'), ('b', 'a'), ('a', 'a'), ('a', 'b')]

14.combinations 功能詳解

combinations(iterable, r=None) 返回 iterable 中長度為 r 的有序排列。預設值 r 為 iterable 的長度。 與 permutations 操作不同的是該函式嚴格按照 iterable 中元素的順序進行排列。

import itertools
print(list(itertools.combinations("abc", r=2)))

# 輸出結果如下
[('a', 'b'), ('a', 'c'), ('b', 'c')]

15.combinations_with_replacement 功能詳解

combinations_with_replacement(iterable, r=None) 返回 iterable 中長度為 r 的有序排列。預設值 r 為 iterable 的長度。 與 combinations 操作不同的是該函式允許每個元素重複出現。

import itertools
print(list(itertools.combinations_with_replacement("abc", r=2)))
#學習中遇到問題沒人解答？小編建立了一個Python學習交流群：153708845
# 輸出結果如下
[('a', 'a'), ('a', 'b'), ('a', 'c'), ('b', 'b'), ('b', 'c'), ('c', 'c')]

itertools 總結

本文總結了 itertools 模組的常規操作，學習並掌握這些極為便利的操作非常有助於提高自己的編碼效率。