迭代器相關(guān)「itertools」模塊使用解讀

0 前言

說到處理循環(huán)，我們習(xí)慣使用for, while等，比如依次打印每個(gè)列表中的字符：

lis?=?['I',?'love',?'python']
for?i?in?lis:
????print(i)
I
love
python

在打印內(nèi)容字節(jié)數(shù)較小時(shí)，全部載入內(nèi)存后，再打印，沒有問題。可是，如果現(xiàn)在有成千上百萬(wàn)條車輛行駛軌跡，叫你分析出其中每個(gè)客戶的出行規(guī)律，堵車情況等，假如是在單機(jī)上處理這件事。

你可能首先要面臨，也可能被你忽視，最后代碼都寫好后，才可能暴露出的一個(gè)問題:outofmemory, 這在實(shí)際項(xiàng)目中經(jīng)常遇到。

這個(gè)問題提醒我們，處理數(shù)據(jù)時(shí)，如何寫出高效利用內(nèi)存的程序，就顯得很重要。今天，我們就來探討如何高效利用內(nèi)存，節(jié)省內(nèi)存同時(shí)還能把事情辦好。

其實(shí)，Python已經(jīng)準(zhǔn)備好一個(gè)模塊專門用來處理這件事，它就是 itertools 模塊，這里面幾個(gè)函數(shù)的功能其實(shí)很好理解。

我不打算籠統(tǒng)的介紹它們所能實(shí)現(xiàn)的功能，而是想分析這些功能背后的實(shí)現(xiàn)代碼，它們?nèi)绾巫龅礁咝Ч?jié)省內(nèi)存的，Python內(nèi)核的貢獻(xiàn)者們又是如何寫出一手漂亮的代碼的，這很有趣，不是嗎？

OK，let's go. Hope you enjoy the journey!

1 拼接元素

itertools 中的chain 函數(shù)實(shí)現(xiàn)元素拼接，原型如下，參數(shù)*表示個(gè)數(shù)可變的參數(shù)

chain(iterables)

應(yīng)用如下：

In?[33]:?list(chain(['I','love'],['python'],['very',?'much']))
Out[33]:?['I',?'love',?'python',?'very',?'much']

哇，不能再好用了，它有點(diǎn)join的味道，但是比join強(qiáng)，它的重點(diǎn)在于參數(shù)都是可迭代的實(shí)例。

那么，chain如何實(shí)現(xiàn)高效節(jié)省內(nèi)存的呢？chain大概的實(shí)現(xiàn)代碼如下：

def?chain(*iterables):
????for?it?in?iterables:
????????for?element?in?it:
????????????yield?element

以上代碼不難理解，chain本質(zhì)返回一個(gè)生成器，所以它實(shí)際上是一次讀入一個(gè)元素到內(nèi)存，所以做到最高效地節(jié)省內(nèi)存。

2 逐個(gè)累積

返回列表的累積匯總值，原型：

accumulate(iterable[, func, *, initial=None])

應(yīng)用如下：

In?[36]:?list(accumulate([1,2,3,4,5,6],lambda?x,y:?x*y))
Out[36]:?[1,?2,?6,?24,?120,?720]

accumulate大概的實(shí)現(xiàn)代碼如下：

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簡(jiǎn)單的yield不同，此處稍微復(fù)雜一點(diǎn)，yield有點(diǎn)像return，所以 yield total那行直接就返回一個(gè)元素，也就是iterable的第一個(gè)元素，因?yàn)槿魏螘r(shí)候這個(gè)函數(shù)返回的第一個(gè)元素就是它的第一個(gè)。又因?yàn)閥ield返回的是一個(gè)generator對(duì)象，比如名字gen，所以next(gen)時(shí)，代碼將會(huì)執(zhí)行到 for element in it:這行，而此時(shí)的迭代器it 已經(jīng)指到iterable的第二個(gè)元素，OK，相信你懂了！

3 漏斗篩選

它是compress 函數(shù)，功能類似于漏斗功能，所以我稱它為漏斗篩選，原型：

compress(data, selectors)

In?[38]:?list(compress('abcdefg',[1,1,0,1]))
Out[38]:?['a',?'b',?'d']

容易看出，compress返回的元素個(gè)數(shù)等于兩個(gè)參數(shù)中較短的列表長(zhǎng)度。

它的大概實(shí)現(xiàn)代碼：

def?compress(data,?selectors):
????return?(d?for?d,?s?in?zip(data,?selectors)?if?s)

這個(gè)函數(shù)非常好用

4 段位篩選

掃描列表，不滿足條件處開始往后保留，原型如下：

dropwhile(predicate, iterable)

應(yīng)用例子：

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]

實(shí)現(xiàn)它的大概代碼如下：

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

掃描列表，只要滿足條件就從可迭代對(duì)象中返回元素，直到不滿足條件為止，原型如下：

takewhile(predicate, iterable)

應(yīng)用例子：

In?[43]:?list(takewhile(lambda?x:?x<5,?[1,4,6,4,1]))
Out[43]:?[1,?4]

實(shí)現(xiàn)它的大概代碼如下：

def?takewhile(predicate,?iterable):
????for?x?in?iterable:
????????if?predicate(x):
????????????yield?x
????????else:
????????????break?#立即返回

6 次品篩選

掃描列表，只要不滿足條件都保留，原型如下：

dropwhile(predicate, iterable)

應(yīng)用例子：

In?[40]:?list(filterfalse(lambda?x:?x%2==0,?[1,2,3,4,5,6]))
Out[40]:?[1,?3,?5]

實(shí)現(xiàn)它的大概代碼如下：

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 必須全部載入內(nèi)存，所以更節(jié)省內(nèi)存的操作islice，原型如下：

islice(iterable, start, stop[, step])

應(yīng)用例子：

In?[41]:?list(islice('abcdefg',1,4,2))
Out[41]:?['b',?'d']

實(shí)現(xiàn)它的大概代碼如下：

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

巧妙利用生成器迭代結(jié)束時(shí)會(huì)拋出異常StopIteration，做一些邊界處理的事情。

8 細(xì)胞分裂

tee函數(shù)類似于我們熟知的細(xì)胞分裂，它能復(fù)制原迭代器n個(gè)，原型如下：

tee(iterable, n=2)

應(yīng)用如下，可以看出復(fù)制出的兩個(gè)迭代器是獨(dú)立的

a?=?tee([1,4,6,4,1],2)
In?[51]:?next(a[0])
Out[51]:?1

In?[52]:?next(a[1])
Out[52]:?1

實(shí)現(xiàn)它的代碼大概如下：

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 實(shí)現(xiàn)內(nèi)部使用一個(gè)隊(duì)列類型deques，起初生成空隊(duì)列，向復(fù)制出來的每個(gè)隊(duì)列中添加元素newval, 同時(shí)yield 當(dāng)前被調(diào)用的mydeque中的最左元素。

9 map變體

starmap可以看做是map的變體，它能更加節(jié)省內(nèi)存，同時(shí)iterable的元素必須也為可迭代對(duì)象，原型如下：

starmap(function, iterable)

應(yīng)用它：

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的實(shí)現(xiàn)細(xì)節(jié)如下：

def?starmap(function,?iterable):
????for?args?in?iterable:
????????yield?function(*args)

10 復(fù)制元素

repeat實(shí)現(xiàn)復(fù)制元素n次，原型如下：

repeat(object[, times])

應(yīng)用如下：

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]]

它的實(shí)現(xiàn)細(xì)節(jié)大概如下：

def?repeat(object,?times=None):
????if?times?is?None:#?如果times不設(shè)置，將一直repeat下去
????????while?True:?
????????????yield?object
????else:
????????for?i?in?range(times):
????????????yield?object

11 笛卡爾積

笛卡爾積實(shí)現(xiàn)的效果同下：

?((x,y)?for?x?in?A?for?y?in?B)

所以，笛卡爾積的實(shí)現(xiàn)效果如下：

In?[68]:?list(product('ABCD',?'xy'))
Out[68]:
[('A',?'x'),
?('A',?'y'),
?('B',?'x'),
?('B',?'y'),
?('C',?'x'),
?('C',?'y'),
?('D',?'x'),
?('D',?'y')]

它的實(shí)現(xiàn)細(xì)節(jié)：

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 加強(qiáng)版zip

組合值。若可迭代對(duì)象的長(zhǎng)度未對(duì)齊，將根據(jù) fillvalue 填充缺失值，注意：迭代持續(xù)到耗光最長(zhǎng)的可迭代對(duì)象，效果如下：

In?[69]:?list(zip_longest('ABCD',?'xy',?fillvalue='-'))
Out[69]:?[('A',?'x'),?('B',?'y'),?('C',?'-'),?('D',?'-')]

它的實(shí)現(xiàn)細(xì)節(jié)：

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，也就是在可迭代對(duì)象的長(zhǎng)度未對(duì)齊時(shí)，根據(jù) fillvalue 填充缺失值。理解上面代碼的關(guān)鍵是迭代器對(duì)象(iter)，next方法的特殊性：

In?[74]:?for?i,?it?in?enumerate([iter([1,2,3]),iter(['x','y'])]):
????...:?????print(next(it))
????#輸出：
????1
????x

結(jié)合這個(gè)提示再理解上面代碼，就不會(huì)吃力。

總結(jié)

Python的itertools模塊提供的節(jié)省內(nèi)存的高效迭代器，里面實(shí)現(xiàn)基本都借助于生成器，所以一方面了解這12個(gè)函數(shù)所實(shí)現(xiàn)的基本功能，同時(shí)也能加深對(duì)生成器(generator)的理解，為我們寫出更加高效、簡(jiǎn)潔、漂亮的代碼打下堅(jiān)實(shí)基礎(chǔ)。