理解zip函式的工作流程

老齊Py發表於2019-03-03
函式

zip函式是Python的內建函式，在拙作《跟老齊學Python：輕鬆入門》有一定的介紹，但是，考慮到那本書屬於Python的入門讀物，並沒有講太深。但是，並不意味著這個函式不能應用的很深入，特別是結合迭代器來理解此函式，會讓人有豁然開朗的感覺。同時，能夠巧妙地解決某些問題。

本文試圖對zip進行深入探討，同時兼顧對迭代器的理解。

一、理解zip函式

看下面的操作。

>>> list(zip([1, 2, 3], [`a`, `b`, `c`]))
[(1, `a`), (2, `b`), (3, `c`)]
複製程式碼

[1, 2, 3] 和 [`a`, `b`, `c`]是列表，所有列表都是可迭代的。這意味著可以一次返回一個元素。

zip函式最終得到一個zip物件——一個迭代器物件。而這個迭代器物件是由若干個元組組成的。

那麼這些元組是如何生成的呢？

對照上面的程式碼，從開始算起。按照Python中的技術習慣，開始的那個元組是用0來計數的，即第0個。

  1. 第0個元組中索引為0的元素，來自於zip函式第0個引數（[1,2,3]）當前指標所指的值，剛開始讀取，那就是1；
  2. 第0個元組中索引為1的元素，來自於zip函式中第1個引數（[`a`, `b`, `c`]）當前指標所指的值，也是剛剛開始讀取，應該是a；於是組成了zip物件的第0個元組(1, `a`)。
  3. 第0個元組中索引為2的元素，來自於zip函式中第2個引數——沒有第二個，上面的引數總共有0、1兩個。那麼這個元組就組建完畢了，也就沒有索引為2的元素。接下來再組建下一個元組，按照這裡的計數順序應該是第1個元組。
  4. zip物件第1個元組中索引為0的元素，來自於zip函式第0個引數當前指標所指的值。注意，因為上次已經度去過1了，這時候對應的值是2。
  5. 同上面的道理，這時候指標所指的值是`b`。於是組成了zip物件第1個元組(2, `b`)。
  6. 如此重複，就得到了最終的結果。

請注意上面的敘述，如果把元組中的組成物件來源概括一句話，那就是：元組中的第i個元素就來自於第i個引數中指標在當前所指的元素物件——請細細品味這句話的含義，後面有大用途。

對於zip函式，上面的過程，貌似“壓縮”一樣，那麼，是否有反過程——解壓縮。例如從上面示例的結果中分別恢復出來原來的兩個物件。

有。一般認為是這這麼做：

>>> result = zip([1, 2, 3], ["a", "b", "c"])
>>> c, v = zip(*result)
>>> print(c, v)
(1, 2, 3) (`a`, `b`, `c`)
複製程式碼

這是什麼原理。

result應用的是一個迭代器物件，不過為了能夠顯示的明白，也可以理解為是[(1, `a`), (2, `b`), (3, `c`)]。接下來使用了zip(*result)，這裡的符號*的作用是收集引數，在函式的引數中，有對此詳細闡述（請參閱《跟老齊學Python：輕鬆入門》）。

>>> def foo(*a): print(a)
...
>>> lst = [(1,2), (3,4)]
>>> foo(*lst)
((1, 2), (3, 4))
>>> foo((1,2), (3,4))
((1, 2), (3, 4))
複製程式碼

仿照這個示例，就能明晰下面兩個操作是等效的。

>>> lst = [(1, `a`), (2, `b`), (3, `c`)]
>>> zip(*lst)
<zip object at 0x104c8fb48>
>>> zip((1, `a`), (2, `b`), (3, `c`))
<zip object at 0x104f27308>
複製程式碼

從返回的物件記憶體編碼也可以看出，兩個是同樣的物件。

既然如此，我們就可以通過理解zip((1, `a`), (2, `b`), (3, `c`))的結果產生過程來理解zip(*lst)了。而前者生成結果的過程前面已經闡述過了，此處不再贅述。

原來，所謂的“解壓縮”和“壓縮”，計算的方法是一樣的。豁然開朗。

其它常用函式

除了zip函式，還有一個內建函式iter，它以可迭代物件為引數，會返回迭代器物件。

>>> iter([1, 2, 3, 4])
<list_iterator object at 0x7fa80af0d898>
複製程式碼

本質上，iter函式呼叫引數的每個元素，然後藉助於__next__函式返回迭代器物件，並把結果整合到一個元組中。

內建函式map是另外一個返回迭代器物件的函式，它以只有一個引數的函式物件為引數，這個函式每次從可迭代物件中取一個元素。

>>> list(map(len, [`abc`, `de`, `fghi`]))
[3, 2, 4]
複製程式碼

map函式的執行原理是：用__iter__函式呼叫第二個引數，並用__next__函式返回執行結果。在上面的例子中，len函式要呼叫後面的列表中的每個元素，並返回一個迭代器物件。

既然迭代器是可迭代的，就可以把zip返回的迭代器物件用到map函式的引數中了。例如，用下面的方式計算兩個列表中對應元素的和。

>>> list(map(sum, zip([1, 2, 3], [4, 5, 6])))
[5, 7, 9]
複製程式碼

迭代器物件有兩個主要功效，一是節省記憶體，而是提高執行效率。

典型問題

有一個列表，由一些正整陣列成，如[1, 2, 3, 4, 5, 6]，寫一個函式，函式的一個引數n表示要將列表中幾個元素劃為一組，假設n=2，則將兩個元素為一組，最終返回[(1, 2), (3, 4), (5, 6)]。

如果用簡單的方式，可以這樣寫此函式：

def naive_grouper(inputs, n):
    num_groups = len(inputs) // n
    return [tuple(inputs[i*n:(i+1)*n]) for i in range(num_groups)]
複製程式碼

測試一下，此函式能夠如我們所願那樣工作。

>>> nums = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
>>> naive_grouper(nums, 2)
[(1, 2), (3, 4), (5, 6), (7, 8), (9, 10)]
複製程式碼

但是，在上面的測試中，所傳入的列表元素個數是比較小的，如果列表元素個數很多，比如有100萬個。這就需要有比較大的記憶體了，否則無法執行運算。

但是，如果這樣執行此程式：

def naive_grouper(inputs, n):
    num_groups = len(inputs) // n
    return [tuple(inputs[i*n:(i+1)*n]) for i in range(num_groups)]
 
for _ in naive_grouper(range(100000000), 10):
    pass
複製程式碼

把上面的程式儲存為檔案naive.py。可以用下面的指令，測量執行程式時所佔用的記憶體空間和耗費的時長。注意，要確保你本地機器的記憶體至少5G。

$ time -f "Memory used (kB): %M
User time (seconds): %U" python3 naive.py
Memory used (kB): 4551872
User time (seconds): 11.04
複製程式碼

注意：Ubuntu系統中，你可能要執行 /usr/bin/time。

把列表或者元素傳入naïve_grouper函式，需要計算機提供4.5GB的記憶體空間，才能執行range(100000000)的迴圈。

如果採用迭代器物件，就會有很大變化了。

def better_grouper(inputs, n):
    iters = [iter(inputs)] * n
    return zip(*iters)
複製程式碼

這個簡短的函式中，內涵還是很豐富的。所以我們要逐行解釋。

表示式[iters(inputs)] * n建立一個迭代器物件，它包含了n個同樣的列表物件。

下面以n=2為例說明。

>>> nums = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
>>> iters = [iter(nums)] * 2
>>> list(id(itr) for itr in iters)  # 記憶體地址是一樣的
[139949748267160, 139949748267160]
複製程式碼

在iters中的兩個迭代器物件是同一個物件——認識到這一點非常重要。

結合前面對zip的理解，zip(*iters)和zip(iter(nums), iter(nums))是一樣的。為了能夠以更直觀的方式進行說明，就可以認為是zip((1, 2, 3, 4, 5, 6, 7, 8, 9, 10), (1, 2, 3, 4, 5, 6, 7, 8, 9, 10))。按照前文所述的zip工作流程，其計算過程如下：

  1. 第0個元組中的第0個元素，來自於第0個引數中指標當前所指物件，是1；
  2. 第0個元組中的第1個元素，來自於第1個引數中指標當前所指物件，特別注意，兩個引數是同一個物件，此時指標所指的是2。
  3. 第0個元組中的第2個元素，來自於第2個引數——沒有。於是乎得到了元組(1, 2)。
  4. 接下來是第1個元組中的第0個元素，來自第0個引數中指標當前所指的物件，還是因為同一個物件的元素，此時指標所指的是3。
  5. 依次類推，得到(3,4)等元組。
>>> nums = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
>>> list(better_grouper(nums, 2))
[(1, 2), (3, 4), (5, 6), (7, 8), (9, 10)]
複製程式碼

上面的函式better_grouper()的優點在於：

  • 不使用內建函式len()，能夠以任何可迭代物件為引數
  • 返回的是可迭代物件，不是列表。這樣更少佔用記憶體

把上述流程儲存為檔案better.py。

def better_grouper(inputs, n):
    iters = [iter(inputs)] * n
    return zip(*iters)
  
for _ in better_grouper(range(100000000), 10):
    pass
複製程式碼

然後使用 time在終端執行。

$ time -f "Memory used (kB): %M
User time (seconds): %U" python3 better.py
Memory used (kB): 7224
User time (seconds): 2.48
複製程式碼

對比前面執行 naive.py，不論在記憶體還是執行時間上，都表現非常優秀。

進一步研究

對於上面的better_grouper函式，深入分析一下，發現它還有問題。它只能分割能夠被列表長度整除的列表中的數字，如果不能整除的話，就會有一些元素被捨棄。例如：

>>> nums = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
>>> list(better_grouper(nums, 4))
[(1, 2, 3, 4), (5, 6, 7, 8)]
複製程式碼

如果要4個元素一組，就會有9和10不能分組了。之所以，還是因為zip函式。

>>> list(zip([1, 2, 3], [`a`, `b`, `c`, `d`]))
[(1, `a`), (2, `b`), (3, `c`)]
複製程式碼

最終得到的zip物件中的元組數量，是引數中長度最小的物件決定。

如果你感覺這樣做不爽，可以使用itertools.zip_longest()，這個函式是以最長的引數為基準，如果有不足的，預設用None填充，當然也可以通過引數fillvalue指定填充物件。

>>> import itertools
>>> x = [1, 2, 3]
>>> y = ["a", "b", "c", "d"]
>>> list(itertools.zip_longest(x, y))
[(1, `a`), (2, `b`), (3, `c`), (None, `d`)]
複製程式碼

那麼，就可以將better_grouper函式中的zip用zip_longest()替代了。

import itertools as it
  
def grouper(inputs, n, fillvalue=None):
    iters = [iter(inputs)] * n
    return it.zip_longest(*iters, fillvalue=fillvalue)
複製程式碼

再跑一下，就是這樣的結果了。

>>> nums = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
>>> print(list(grouper(nums, 4)))
[(1, 2, 3, 4), (5, 6, 7, 8), (9, 10, None, None)]
複製程式碼

《跟老齊學Python：輕鬆入門》

相關文章