Python的美麗在于它的簡潔性。

不僅因為Python的語法優(yōu)雅，還因為它有許多設計良好的內(nèi)置模塊，能夠高效地實現(xiàn)常見功能。

itertools模塊就是一個很好的例子，它為我們提供了許多強大的工具，可以在更短的代碼中操作Python的可迭代對象。

用更少的代碼實現(xiàn)更多的功能，這就是你可以從itertools模塊中獲得的好處。讓我們從本文中了解一下。

1、itertools.product(): 避免嵌套循環(huán)的巧妙方法

當程序變得越來越復雜時，你可能需要編寫嵌套循環(huán)。同時，你的Python代碼將變得丑陋和難以閱讀：

list_a = [1, 2020, 70]
list_b = [2, 4, 7, 2000]
list_c = [3, 70, 7]


for a in list_a:
    for b in list_b:
        for c in list_c:
            if a + b + c == 2077:
                print(a, b, c)
# 70 2000 7

如何使上述代碼再次具有 Python 風格？

那 itertools.product() 函數(shù)就是你的朋友：

from itertools import product


list_a = [1, 2020, 70]
list_b = [2, 4, 7, 2000]
list_c = [3, 70, 7]


for a, b, c in product(list_a, list_b, list_c):
    if a + b + c == 2077:
        print(a, b, c)
# 70 2000 7

如上所示，它返回輸入可迭代對象的笛卡爾積，幫助我們將三個嵌套的for循環(huán)合并為一個。

2、itertools.compress(): 過濾數(shù)據(jù)的便捷方式

我們可以通過一個或多個循環(huán)來篩選列表中的項。

但有時候，我們可能不需要編寫任何循環(huán)。因為有一個名為itertools.compress()的函數(shù)。

itertools.compress()函數(shù)返回一個迭代器，根據(jù)相應的布爾掩碼對可迭代對象進行過濾。

例如，以下代碼使用itertools.compress()函數(shù)選擇出真正的領導者：

import itertools
leaders = ['Yang', 'Elon', 'Tim', 'Tom', 'Mark']
selector = [1, 1, 0, 0, 0]
print(list(itertools.compress(leaders, selector)))
# ['Yang', 'Elon']

第二個參數(shù)selector作為掩碼起作用，我們也可以這樣定義它：

selector = [True, True, False, False, False]

3、itertools.groupby(): 對可迭代對象進行分組

itertools.groupby()函數(shù)是一種方便的方式，用于將可迭代對象中相鄰的重復項進行分組。

例如，我們可以將一個長字符串進行分組，如下所示：

from itertools import groupby


for key, group in groupby('LinnuxmiMi'):
    print(key, list(group))

此外，我們可以利用它的第二個參數(shù)告訴groupby()函數(shù)如何確定兩個項是否相同：

from itertools import groupby


for key, group in groupby('LinnuxmiMi', lambda x: x.upper()):
    print(key, list(group))

4、itertools.combinations(): 從可迭代對象中獲取給定長度的所有組合

對于初學者來說，編寫一個無 bug 的函數(shù)來獲取列表的所有可能組合可能需要一些時間。

事實上，如果她了解 itertools.combinations() 函數(shù)，她可以很容易地實現(xiàn)：

import itertools


author = ['L', 'i', 'n', 'u', 'x']


result = itertools.combinations(author, 2)


for a in result:
    print(a)

如上所示，itertools.combinations()函數(shù)有兩個參數(shù)，一個是原始可迭代對象，另一個是函數(shù)生成的子序列的長度。

5、itertools.permutations(): 從可迭代對象中獲取給定長度的所有排列

既然有一個函數(shù)可以獲取所有組合，當然也有另一個名為itertools.permutations的函數(shù)來獲取所有可能的排列：

import itertools


author = ['Y', 'a', 'n', 'g']


result = itertools.permutations(author, 2)


for x in result:
    print(x)


# ('Y', 'a')
# ('Y', 'n')
# ('Y', 'g')
# ('a', 'Y')
# ('a', 'n')
# ('a', 'g')
# ('n', 'Y')
# ('n', 'a')
# ('n', 'g')
# ('g', 'Y')
# ('g', 'a')
# ('g', 'n')

如上所示，itertools.permutations()函數(shù)的用法與itertools.combinations()類似。唯一的區(qū)別在于它們的結(jié)果。

6、itertools.accumulate(): 從可迭代對象生成累積的項

基于可迭代對象獲取一系列累積值是一種常見需求。借助itertools.accumulate()函數(shù)的幫助，我們無需編寫任何循環(huán)即可實現(xiàn)。

import itertools
import operator


nums = [1, 2, 3, 4, 5]
print(list(itertools.accumulate(nums, operator.mul)))
# [1, 2, 6, 24, 120]

如果我們不想使用operator.mul，上述程序可以改寫如下：

import itertools


nums = [1, 2, 3, 4, 5]
print(list(itertools.accumulate(nums, lambda a, b: a * b)))
# [1, 2, 6, 24, 120]

7、itertools.repeat(), itertools.cycle(), itertools.count(): 生成無限迭代對象

在某些情況下，我們需要獲得無限迭代。有 3 個有用的功能：

itertools.repeat（）：重復生成相同的項

例如，我們可以得到三個相同的“Yang”，如下所示：

import itertools
print(list(itertools.repeat('Yang', 3)))
# ['Yang', 'Yang', 'Yang']

itertools.cycle(): 通過循環(huán)獲得無限迭代器

itertools.cycle函數(shù)將不會停止，直到我們跳出循環(huán)：

import itertools


count = 0


for c in itertools.cycle('Yang'):
    if count >= 12:
        break
    else:
        print(c, end=',')
        count += 1
# Y,a,n,g,Y,a,n,g,Y,a,n,g,

itertools.count(): 生成一個無限的數(shù)字序列 如果我們只需要數(shù)字，可以使用itertools.count函數(shù)：

import itertools


for i in itertools.count(0, 2):
    if i == 20:
        break
    else:
        print(i, end=" ")
# 0 2 4 6 8 10 12 14 16 18

如上所示，它的第一個參數(shù)是起始數(shù)字，第二個參數(shù)是步長。

8、itertools.pairwise(): 輕松獲取成對的元組

自從Python 3.10版本開始，itertools模塊新增了一個名為pairwise的函數(shù)。它是一個簡潔而方便的工具，用于從可迭代對象中生成連續(xù)的重疊對。

import itertools


letters = ['a', 'b', 'c', 'd', 'e']


result = itertools.pairwise(letters)


print(list(result))
# [('a', 'b'), ('b', 'c'), ('c', 'd'), ('d', 'e')]

9、itertools.takewhile(): 以不同的方式過濾元素

itertools.takewhile()返回一個迭代器，只要給定的謂詞函數(shù)評估為True，就會從可迭代對象中生成元素。

import itertools


nums = [1, 61, 7, 9, 2077]


print(list(itertools.takewhile(lambda x: x < 100, nums)))
# [1, 61, 7, 9]

該函數(shù)與內(nèi)置的filter()函數(shù)不同。

filter函數(shù)將遍歷整個列表：

nums = [1, 61, 7, 9, 2077]


print(list(filter(lambda x: x < 10, nums)))
# [1, 7, 9]

然而，itertools.takewhile函數(shù)如其名稱所示，當評估函數(shù)為False時會停止迭代：

import itertools


nums = [1, 61, 7, 9, 2077]


print(list(itertools.takewhile(lambda x: x < 10, nums)))
# [1]

10、itertools.dropwhile(): itertools.takewhile的反向操作

這個函數(shù)似乎是前面那個函數(shù)的相反思路。

itertools.takewhile()函數(shù)在謂詞函數(shù)為True時返回可迭代對象的元素，而itertools.dropwhile()函數(shù)在謂詞函數(shù)為True時丟棄可迭代對象的元素，然后返回剩下的元素。

import itertools


nums = [1, 61, 7, 9, 2077]


print(list(itertools.dropwhile(lambda x: x < 100, nums)))
# [2077]

審核編輯：湯梓紅