Python 迭代器模組 itertools 簡介

張無忌發表於2016-06-07

Python提供了一個非常棒的模組用於建立自定義的迭代器,這個模組就是 itertools。itertools 提供的工具相當高效且節省記憶體。使用這些工具,你將能夠建立自己定製的迭代器用於高效率的迴圈。這一章,我們將一起看一看這些工具的應用例項以便理解並應用到自己的程式設計中去。

讓我們先從幾個無限迭代器的例子開始吧!


無限迭代器

itertools 包自帶了三個可以無限迭代的迭代器。這意味著,當你使用他們時,你要知道你需要的到底是最終會停止的迭代器,還是需要無限地迭代下去。

這些無限迭代器在生成數字或者在長度未知的可迭代物件(iterables)中迴圈時相當有用。下面我們開始認識這些有趣的可迭代物件!

count(初值=0, 步長=1)

count 迭代器會返回從傳入的起始引數開始的均勻間隔的數值。count 也可以接收指定的步長引數。我們來看一個簡單的例子:

這裡我們先從 itertools 匯入 count,然後建立一個 for 迴圈。迴圈中加入了條件檢查,當迭代器值大於 20 時跳出迴圈,否則將列印出迭代器的當前值。你應該注意到了,輸出結果從 10 開始,與我們傳入 count 的起始值是一致的。

另一種控制無限迭代器輸出的方式是使用 itertools 的子模組 islice。使用方法如下:

這裡,我們先匯入了 islice,然後遍歷 count,從 10 開始,輸出 5 個元素後結束。你大概猜到了,islice 的第二個引數控制何時停止迭代。但其含義並不是”達到數字 5 時停止“,而是”當迭代了 5 次之後停止“。

cycle(可迭代物件)

itertools 中的 cycle 迭代器允許你建立一個能在一組值間無限迴圈的迭代器。下面我們傳入一個 3 個字母的字串,看看將會發生什麼:

這裡我們建立了一個 for 迴圈,使其在三個字母 XYZ 間無限迴圈。當然,我們並不真地想要永遠迴圈下去,所以我們新增了一個簡單的計數器來跳出迴圈。

你也可以用 Python 內建的 next 函式對 itertools 建立的迭代器進行迴圈:

上面的程式碼中,我們建立一個簡單的多邊形組成的列表,然後傳入 cycle。我們用一個變數儲存新建的迭代器,然後將這個變數傳入 next 函式。每一次呼叫 next 都將返回迭代器中的下一個值。由於迭代器是無限的,我們可以一直呼叫 next 而永遠不會到盡頭。

repeat(物件[, 次數])

repeat 迭代器會一遍遍地返回傳入的物件直至永遠,除非你設定了 times 引數。這與 cycle 非常相像,除了一點,repeat 不會在多個值之間迴圈。我們看一個簡單的例子:

這裡,我們先匯入 repeat,然後設定其重複五次數字 5。接著,我們連續六次呼叫 next 函式,觀察它是否工作正常。當執行這段程式碼,將引發 StopIteration,因為最後一次呼叫 next 函式,迭代器中的值已經全部返回了。


可終止的迭代器

多數用 itertools 建立的迭代器並不是無限的。這一節,我們將瞭解 itertools 中的有限迭代器。為了對於輸出的可讀的結果,我們將使用 Python 內建列表儲存。 如果不使用列表,你將只能列印出 itertools 物件。

accumulate(可迭代物件[, 函式])

accumulate 迭代器將返回累計求和結果,或者傳入兩個引數的話,由傳入的函式累積計算的結果。預設設定為相加,我們趕快試一試吧:

這裡,我們 匯入了 accumulate,然後傳入 10 個數字,0-9。迭代器將傳入數字依次累加,所以第一個是 0 ,第二個是 0+1, 第三個是 1+2,如此下去。現在我們匯入 operator 模組,然後新增進去:

這裡我們傳入了數字 1-4 到 accumulate 迭代器中。我們還傳入了一個函式:operator.mul,這個函式將接收的引數相乘。所以每一次迭代,迭代器將以乘法代替除法(1×1=1, 1×2=2, 2×3=6, 以此類推)。

accumulate 的文件中給出了其他一些有趣的例子,例如貸款分期償還,混沌遞推關係等。這絕對值得你花時間去看一看。

chain(*可迭代物件)

chain 迭代器能夠將多個可迭代物件合併成一個更長的可迭代物件。實際上,我參與的一個專案中最近就需要這一功能。我有一個列表,裡面已經包含一些元素,接著想把另外兩個列表新增到最初那個列表中。注意,我們想新增的是兩個列表的元素。最初,我是這樣做的:

這並不是我想要的。itertools 模組提供一個優雅得多的方法用chain 來合併這些列表:

許多聰明的讀者可能想到了,實際上不使用 itertools,也有其他方法能夠實現這一要求。你可以這樣做:

這些方法當然都是可行的。在我知道 chain 之前,我可能會這樣做,但我個人認為這個例子中, chain 更為優雅,也更容易理解。

chain.from_iterable(可迭代物件)

你也可以用 chain 的一個方法,叫做 from_iterable。這個方法與直接用 chain 有些細微的差別。不同於直接傳入一系列可迭代物件,你必須傳入一個巢狀的列表。我們這就來看一看:

這裡我們跟之前一樣先匯入 chain。我們嘗試傳入兩個列表,但卻出現了 TypeError!為了修正這個錯誤,我們將 cmd 和 numbers 放入一個列表中,將這個巢狀的列表傳入 from_iterable。雖然有點細微差別,但也是很方便使用的。

compress(資料, 選擇器)

子模組 compress 在用一個可迭代物件過濾另一個可迭代物件時十分有用。這是通過將作為選擇器的可迭代物件取為布林值列表來實現的。下面是它的實現方式:

這個例子中,我們有七個字母和五個布林值構成的列表。我們它們傳入 compress 函式。compress 函式將交替檢視兩個可迭代物件。先檢查第一個可迭代物件,然後是第二個,如果第二個的元素為 True,則保留第一個對應元素,如果為 False,則丟棄第一個對應元素。據此,再看看上面的例子,你會發現第一、第三和第四個位置元素為 True,對應的與第一個物件中的 A, C 和 D。

dropwhile(斷言, 可迭代物件)

itertools 還提供了一個小清新的迭代器 dropwhile。只要過濾器判定是  True,dropwhile 迭代器就會排除這些元素。因此,在出現斷言為 False 之前,你不會看到任何輸出結果。這可能導致啟動時間非常長,這點應當注意。

我們看一個來自於 Python 文件的例子:

這裡,我們先匯入 dropwhile,然後傳入一個簡單的 lambda 宣告的函式,如果 x 小於5,lambda 將返回 True ,否則返回 False。dropwhile 將遍歷整個列表,然後將每個元素傳入 lambda。如果 lambda 返回 True 則捨棄該值。一旦遇到了數字 6,lambda 將返回 False,因此,我們將保留數字 6 以及之後餘下的元素。

當我瞭解更多之後發現,用一般的函式比 lambda 函式更為有用。所以我們嘗試建立一個函式,如果引數大於 5,函式返回 True。

這裡,我們在 Python 直譯器中建立了一個簡單的函式作為斷言或過濾器。如果我們傳入的值為 True,這些值都會被捨棄。一旦我們遇到了一個小於 5 的值,那麼該值,包括其後餘下的全部值都將被保留,正如你在例子中看到的。

filterfalse(斷言, 可迭代物件)

itertools 中的 filterfalse 函式與 dropwhile 非常類似。只是,filterfalse 返回斷言為 False 的那些值,而不是捨棄斷言為 True 的值。以我們上一節建立的函式為例來闡釋:

這裡,我們向 filter 傳入了我們建立的函式和一個整數元素的列表。如果整數值小於 5,這個整數將保留下來,否則將被捨棄。注意到,這裡的結果僅有 1, 2 和 3。與 dropwhile 不同,filterfalse將對全部的值進行條件判斷。

groupby(可迭代物件, =None)

groupby 迭代器會從傳入的可迭代物件返回連續的鍵和組。不借助例子可能很難理解這一點,所以我們還是看例子。將下面的程式碼輸入直譯器或者存為一個檔案:

這裡,我們先匯入 groupby,然後建立一個元組構成的列表。接著,我們試圖對資料排序使其輸出時更有意義,這也讓 groubby 能夠正確地對元素分組。下一步,我們遍歷 groupby 返回的含有鍵和組的迭代器。接著我們遍歷組,並且列印出它的內容。如果你執行這段程式碼,你應該會看到下面的結果。

試試改動一下上面的程式碼,將傳入的 sorted_vehicles 替換成 vehicles。你很快就會明白為什麼你要在傳入 groupby 前對資料排序。

islice(可迭代變數, 起始值, 終止值[, 步長])

我們實際上在 count 一節的時候已經提到了 islice,這裡我們將更深入地探討這一函式。islice 是一個返回可迭代物件選定元素的迭代器。這種表述有點模糊。簡而言之,islice 所做的是利用可迭代物件的索引實現切片,然後以迭代器的形式返回所選元素。實際上 islice 有兩種實現方式,一種是 itertools.islice(iterable, stop),還有一種更符合 Python 慣例的形式:islice(iterable, start, stop[, step])。

我們來看看第一種形式,觀察它是如何工作的:

在以上的程式碼中,我們傳入了一個 6 個字元組成的字串以及終止引數 4 給 isilce。這表示 islice 返回的迭代器將包含傳入字串的前 4 個字元。為了驗證,我們連續四次呼叫 next 函式。Python 能夠智慧地識別是否只有 2 個傳入引數,若是,則第二個引數就取為終止引數。

我們再試試傳入三個引數來驗證是否可以同時傳入起始值和終止值。count 工具可以幫助我們解釋這一概念:

這裡我們呼叫了 count,並指定 islice 從 3 開始,到 15 時結束。這跟我們用 slice 效果是一樣的,不同之處在於傳入的是迭代器,返回的是新的迭代器。

starmap(函式, 可迭代物件)

starmap 工具能夠建立一個用傳入的函式和可迭代物件計算的迭代器。如文件中所言,“map() 和 starmap() 的區別正如 function(a,b) 和 function(*c) 的區別。”

下面來看一個簡單的例子:

這裡,我們先建立了一個接受兩個引數的求和函式。接下來我們建立了一個 for 迴圈,並呼叫 starmap 函式,starmap 函式的第一個傳入引數是我們建立的求和函式,第二個傳入引數是元組構成的列表。starmap 函式接著會將每一個元組傳入求和函式,然後將結果返回到迭代器,正如我們列印出的那樣。

takewhile(斷言, 可迭代物件)

takewhile 模組與我們之前介紹的 dropwhile 迭代器剛好相反。takewhile 所建立的迭代器,一旦可迭代物件的斷言或過濾器結果為 True 就返回元素。試一試下面的例子,看看它是如何工作的:

這裡,執行 takewhile 時傳入了一個 lambda 函式和一個列表。輸出的僅是可迭代物件的前兩個整數元素。因為 1 和 4 都小於 5,而 6 大於5,一旦 takewhile 遇到 6, 判定條件結果將是 False,可迭代物件餘下的元素將會被忽略。

tee(可迭代物件, n=2)

tee 工具能夠從一個可迭代物件建立 n 個迭代器。這意味著你能夠用一個可迭代物件建立多個迭代器。下面這段程式碼能大體解釋它是如何工作的:

這裡我們建立了一個 5 個字母組成的字串,然後傳入 tee。由於 tee 的預設引數是 2,我們用兩個變數接收 tee 返回的兩個迭代器。最後,我們分別遍歷了這兩個迭代器,並列印出它們的內容。正如你所見的,它們的內容是相同的。

zip_longest(*可迭代物件, 填充值=None)

zip_longest 迭代器可以用於將兩個可迭代物件配對。如果可迭代物件的長度不同,也可以傳入填充值。我們來看一個來自該函式文件的一個很初級的例子:

這段程式碼裡,我們先匯入 zip_longest,然後傳入兩個需要配對的字串。你會發現,第一個字串有 4 個字元,而第二個只有 2 個字元。我們還傳入了填充值”BLANK“。當遍歷並列印的時候,會看到返回的是元組。前兩個元組是分別來自兩個字串第一個和第二個字母的組合。後兩個則插入了填充值。

需要注意的是,如果傳入 zip_longest 的可迭代物件是無限的,那麼你應該用類似 islice 的工具在外部處理函式,以便控制呼叫的次數。


組合生成器

itertools 庫提供了四個可用於生成資料排列和組合的迭代器。這一章將介紹這些有趣的迭代器。

combinations(可迭代物件, r)

如果你需要生成資料的組合,Python 已經為你準備了 itertools.combinations。combinations 能夠讓你通過有一定長度的可迭代物件建立一個迭代器。請看:

當你執行這段程式碼時會發現,combinations 返回組合是元組。為了使輸出結果更具可讀性,可以遍歷迭代器將元組組合成字串:

現在容易看到全部的組合結果。注意,combinations 函式得到的組合是依循詞典順序的,因而如果傳入的可迭代物件已經排序,那麼你得到的組合元組也將是按序的。還有一點值得注意的是,combinations 函式不會將每個單獨的元素自己跟自己組合。

combinations_with_replacement(可迭代物件, r)

combinationa_with_replacement 迭代器與 combinations 非常類似。唯一的區別是,它會建立元素自己與自己的組合。我們還是用上一節的例子來闡釋:

正如你所見,我們這裡多了四個輸出:WW, XX, YY 和 ZZ。

product(*可迭代物件, 重複=1)

itertools 包還提供一個小清新的函式用於根據一系列輸入的可迭代物件計算笛卡爾積。沒錯,這就是 product 函式。我們來瞧瞧它是如何工作的!

這裡,我們先匯入 product,然後建立一個元組構成的列表,用變數 arrays 儲存。接下來,我們將這些陣列傳入 product 函式,然後呼叫。你應該注意到了,我們呼叫時使用了 *arrays。這會使傳入的列表被拆解,依次應用於 product 函式。這意味著你傳入的將是三個變數而不是一個。如果有興趣,你可以試試呼叫帶星號的 arrays,看看會發生什麼。

permutations(可迭代變數, r)

itertools 的子模組 permutation 能夠返回由傳入的可迭代物件的元素構成的長度為 r 的排列。與 combination 函式類似,permutation 將以詞典序輸出。請看:

容易發現,輸出結果比 combinations 的輸出結果長一些。當使用時,permutation 函式將遍歷字串的所有排列,但當輸入元素是唯一的時候,不會讓這個元素自己和自己排列。


總結

itertools 是建立迭代器時十分通用的工具集。你可以用這些單獨用這些工具,也可以組合起來使用來建立自定義的迭代器。Python 官方文件提供了許多很棒的例項,你可以參考這些例子發掘這個有用的庫的其他用法。


相關閱讀

相關文章