Python高效編程之itertools模組詳解
- 2019 年 11 月 7 日
- 筆記
0 前言
說到處理循環，我們習慣使用for, while等，比如依次列印每個列表中的字元：
lis = [ I , love , python ] for i in lis: print(i) I love python
在列印內容位元組數較小時，全部載入記憶體後，再列印，沒有問題。可是，如果現在有成千上百萬條車輛行駛軌跡，叫你分析出其中每個客戶的出行規律，堵車情況等，假如是在單機上處理這件事。
你可能首先要面臨，也可能被你忽視，最後程式碼都寫好後，才可能暴露出的一個問題:outofmemory, 這在實際項目中經常遇到。
這個問題提醒我們，處理數據時，如何寫出高效利用記憶體的程式，就顯得很重要。今天，我們就來探討如何高效利用記憶體，節省記憶體同時還能把事情辦好。
其實，Python已經準備好一個模組專門用來處理這件事，它就是 itertools 模組，這裡面幾個函數的功能其實很好理解。
我不打算籠統的介紹它們所能實現的功能，而是想分析這些功能背後的實現程式碼，它們如何做到高效節省記憶體的，Python內核的貢獻者們又是如何寫出一手漂亮的程式碼的，這很有趣，不是嗎？
OK，let s go. Hope you enjoy the journey!
1 拼接元素
itertools 中的chain 函數實現元素拼接，原型如下，參數*表示個數可變的參數
chain(iterables)
應用如下：
In [33]: list(chain([ I , love ],[ python ],[ very , much ])) Out[33]: [ I , love , python , very , much ]
哇，不能再好用了，它有點join的味道，但是比join強，它的重點在於參數都是可迭代的實例。
那麼，chain如何實現高效節省記憶體的呢？chain大概的實現程式碼如下：
def chain(*iterables): for it in iterables: for element in it: yield element
以上程式碼不難理解，chain本質返回一個生成器，所以它實際上是一次讀入一個元素到記憶體，所以做到最高效地節省記憶體。
2 逐個累積
返回列表的累積匯總值，原型：
accumulate(iterable[, func, *, initial=None])
應用如下：
In [36]: list(accumulate([1,2,3,4,5,6],lambda x,y: x*y)) Out[36]: [1, 2, 6, 24, 120, 720]
accumulate大概的實現程式碼如下：
def accumulate(iterable, func=operator.add, *, initial=None): it = iter(iterable) total = initial if initial is None: try: total = next(it) except StopIteration: return yield total for element in it: total = func(total, element) yield total
以上程式碼，你還好嗎？與chain簡單的yield不同，此處稍微複雜一點，yield有點像return，所以 yield total那行直接就返回一個元素，也就是iterable的第一個元素，因為任何時候這個函數返回的第一個元素就是它的第一個。又因為yield返回的是一個generator對象，比如名字gen，所以next(gen)時，程式碼將會執行到 for element in it:這行，而此時的迭代器it 已經指到iterable的第二個元素，OK，相信你懂了！
3 漏斗篩選
它是compress 函數，功能類似於漏斗功能，所以我稱它為漏斗篩選，原型：
compress(data, selectors)
In [38]: list(compress( abcdefg ,[1,1,0,1])) Out[38]: [ a , b , d ]
容易看出，compress返回的元素個數等於兩個參數中較短的列表長度。
它的大概實現程式碼：
def compress(data, selectors): return (d for d, s in zip(data, selectors) if s)
這個函數非常好用
4 段位篩選
掃描列表，不滿足條件處開始往後保留，原型如下：
dropwhile(predicate, iterable)
應用例子：
In [39]: list(dropwhile(lambda x: x<3,[1,0,2,4,1,1,3,5,-5])) Out[39]: [4, 1, 1, 3, 5, -5]
實現它的大概程式碼如下：
def dropwhile(predicate, iterable): iterable = iter(iterable) for x in iterable: if not predicate(x): yield x break for x in iterable: yield x
5 段位篩選2
掃描列表，只要滿足條件就從可迭代對象中返回元素，直到不滿足條件為止，原型如下：
takewhile(predicate, iterable)
應用例子：
In [43]: list(takewhile(lambda x: x<5, [1,4,6,4,1])) Out[43]: [1, 4]
實現它的大概程式碼如下：
def takewhile(predicate, iterable): for x in iterable: if predicate(x): yield x else: break #立即返回
6 次品篩選
掃描列表，只要不滿足條件都保留，原型如下：
dropwhile(predicate, iterable)
應用例子：
In [40]: list(filterfalse(lambda x: x%2==0, [1,2,3,4,5,6])) Out[40]: [1, 3, 5]
實現它的大概程式碼如下：
def dropwhile(predicate, iterable): iterable = iter(iterable) for x in iterable: if not predicate(x): yield x break for x in iterable: yield x
7 切片篩選
Python中的普通切片操作，比如：
lis = [1,3,2,1] lis[:1]
它們的缺陷還是lis 必須全部載入記憶體，所以更節省記憶體的操作islice，原型如下：
islice(iterable, start, stop[, step])
應用例子：
In [41]: list(islice( abcdefg ,1,4,2)) Out[41]: [ b , d ]
實現它的大概程式碼如下：
def islice(iterable, *args): s = slice(*args) start, stop, step = s.start or 0, s.stop or sys.maxsize, s.step or 1 it = iter(range(start, stop, step)) try: nexti = next(it) except StopIteration: for i, element in zip(range(start), iterable): pass return try: for i, element in enumerate(iterable): if i == nexti: yield element nexti = next(it) except StopIteration: for i, element in zip(range(i + 1, stop), iterable): pass
巧妙利用生成器迭代結束時會拋出異常StopIteration，做一些邊界處理的事情。
8 細胞分裂
tee函數類似於我們熟知的細胞分裂，它能複製原迭代器n個，原型如下：
tee(iterable, n=2)
應用如下，可以看出複製出的兩個迭代器是獨立的
a = tee([1,4,6,4,1],2) In [51]: next(a[0]) Out[51]: 1 In [52]: next(a[1]) Out[52]: 1
實現它的程式碼大概如下：
def tee(iterable, n=2): it = iter(iterable) deques = [collections.deque() for i in range(n)] def gen(mydeque): while True: if not mydeque: try: newval = next(it) except StopIteration: return for d in deques: d.append(newval) yield mydeque.popleft() return tuple(gen(d) for d in deques)
tee 實現內部使用一個隊列類型deques，起初生成空隊列，向複製出來的每個隊列中添加元素newval, 同時yield 當前被調用的mydeque中的最左元素。
9 map變體
starmap可以看做是map的變體，它能更加節省記憶體，同時iterable的元素必須也為可迭代對象，原型如下：
starmap(function, iterable)
應用它：
In [63]: list(starmap(lambda x,y: str(x)+ - +str(y), [( a ,1),( b ,2),( c ,3)])) Out[63]: [ a-1 , b-2 , c-3 ]
starmap的實現細節如下：
def starmap(function, iterable): for args in iterable: yield function(*args)
10 複製元素
repeat實現複製元素n次，原型如下：
repeat(object[, times])
應用如下：
In [66]: list(repeat(6,3)) Out[66]: [6, 6, 6] In [67]: list(repeat([1,2,3],2)) Out[67]: [[1, 2, 3], [1, 2, 3]]
它的實現細節大概如下：
def repeat(object, times=None): if times is None:# 如果times不設置，將一直repeat下去 while True: yield object else: for i in range(times): yield object
11 笛卡爾積
笛卡爾積實現的效果同下：
((x,y) for x in A for y in B)
所以，笛卡爾積的實現效果如下：
In [68]: list(product( ABCD , xy )) Out[68]: [( A , x ), ( A , y ), ( B , x ), ( B , y ), ( C , x ), ( C , y ), ( D , x ), ( D , y )]
它的實現細節：
def product(*args, repeat=1): pools = [tuple(pool) for pool in args] * repeat result = [[]] for pool in pools: result = [x+[y] for x in result for y in pool] for prod in result: yield tuple(prod)
12 加強版zip
組合值。若可迭代對象的長度未對齊，將根據 fillvalue 填充缺失值，注意：迭代持續到耗光最長的可迭代對象，效果如下：
In [69]: list(zip_longest( ABCD , xy , fillvalue= - )) Out[69]: [( A , x ), ( B , y ), ( C , - ), ( D , - )]
它的實現細節：
def zip_longest(*args, fillvalue=None): iterators = [iter(it) for it in args] num_active = len(iterators) if not num_active: return while True: values = [] for i, it in enumerate(iterators): try: value = next(it) except StopIteration: num_active -= 1 if not num_active: return iterators[i] = repeat(fillvalue) value = fillvalue values.append(value) yield tuple(values)
它裡面使用repeat，也就是在可迭代對象的長度未對齊時，根據 fillvalue 填充缺失值。理解上面程式碼的關鍵是迭代器對象(iter)，next方法的特殊性：
In [74]: for i, it in enumerate([iter([1,2,3]),iter([ x , y ])]): ...: print(next(it)) #輸出： 1 x
結合這個提示再理解上面程式碼，就不會吃力。
總結
Python的itertools模組提供的節省記憶體的高效迭代器，裡面實現基本都藉助於生成器，所以一方面了解這12個函數所實現的基本功能，同時也能加深對生成器(generator)的理解，為我們寫出更加高效、簡潔、漂亮的程式碼打下堅實基礎。
