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):  
    iteriterable = 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):  
    iteriterable = 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, sstep = s.start or 0, s.stop or sys.maxsize, s.step or 1  
    it = iter(range(start, stop, step))  
    try:  
        nextnexti = 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  
                nextnexti = 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方法的特殊性：I

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

結合這個提示再理解上面代碼，就不會吃力。

總結

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