BG75 itertools加速你的Python程序

如果你聽到有一個人在抱怨Python，那麼他十有八九正在承受Python程序運行速度慢的煎熬。

如果有需要實現任意多步的循環，那麼“無窮迭代器”會很有幫助；如果想對序列數據進行一些常規操作，那麼“有限迭代器”會比較有用；如果需要實現一些排列組合的需求，那麼“排列組合迭代器”會是不二的選擇。除了介紹這三種迭代器的常用功能，下文還會對每一類迭代器與自己手動實現相同功能的程序進行性能的對比。可以先劇透一下結論，用過itertools提供的迭代器之後，你再也不想自己實現這些功能了。

1. 無窮迭代器

所謂“無窮迭代器”是指，可以從迭代器中取值任意次數。如果直接對這樣的迭代器進行遍歷，而不進行跳出操作，那麼將會是一個死循環。

(1). count()

如果需要生成一個任意長度的等差數列，count()函數是一個比range()更好的選擇，因為它支持浮點數步長。下面的例子可以說明這一點。

<code>import itertools as it

it_count = it.count(12, 2.5)
print('type(it_count)={}'.format(type(it_count)))

ret = []
for i in it_count:
    ret.append(i)
    # 如果不加這一行就是死循環
    if i > 20:
        break
print(ret)/<code>

Output:

<code>type(it_count)=<class>
[12, 14.5, 17.0, 19.5, 22.0]/<class>/<code>

(2). cycle()

如果需要循環獲取一個有效序列中的每一個元素，cycle()會很有幫助。

<code>it_cycle = it.cycle([1,2,3,4])
print('type(it_cycle)={}'.format(type(it_cycle)))

ret = []
for i in it_cycle:
    ret.append(i)
    if len(ret) >= 10:
        break
print(ret)/<code>

Output:

<code>type(it_cycle)=<class>
[1, 2, 3, 4, 1, 2, 3, 4, 1, 2]/<class>/<code>

(3). repeat()

如果需要重複一個數值任意次數，可以使用repeat()函數；如果需要重複有限次，可以將次數作為第二個參數傳入。

<code>it_repeat = it.repeat(2)
print('type(it_repeat)={}'.format(type(it_repeat)))

ret = []
for i in it_repeat:
    ret.append(i)
    if len(ret) >= 10:
        break
print(ret)/<code>

Output:

<code>type(it_repeat)=<class>
[2, 2, 2, 2, 2, 2, 2, 2, 2, 2]/<class>/<code>

(4). 性能對比

要說這三個函數所提供的功能更加高效，那總不能空開說白話，需要給出實驗數據才有說服力。我們先實現一個與count()功能一致的函數。

<code>def my_count(n, step):
    """使用python代碼實現itertools.count的功能
    """
    while True:
        yield n
        n += step/<code>

為了便於測試，我們還需要實現一個測試函數，用於完成正常的功能調用。

<code>def test_count(it_count):
    """用於測試count生成器的性能
    """
    for i in it_count: 
        if i > 20: 
            break/<code>

藉助notebook的魔術方法，我們測試一下itertools提供的count()和自己動手實現的my_count()在相同參數下的運行耗時。

<code># 測試itertools提供的函數    
it_mycount = my_count(12, 2.5)
%timeit test_count(it_count)
it_count = it.count(12, 2.5)
%timeit test_count(it_mycount)/<code>

Output:

<code>177 ns ± 2.3 ns per loop (mean ± std. dev. of 7 runs, 10000000 loops each)
220 ns ± 2.57 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)/<code>

評測數據說明，自己實現的my_count()至少比itertools提供的count()慢了16%。結論也就顯而易見了，能用已有的，千萬別瞎動手。

2. 有限迭代器

相比於“無窮迭代器”，更常用的還是“有限迭代器”。在日常開發的代碼中使用這些接口，會比自己實現省事得多。

(1). accumulate()

累加，這是一個常用的數學操作，自己實現起來也不復雜，但在這裡，我們還是先看看itertools的接口。

<code>a = list(range(10))

it_acc = it.accumulate(a)
print('type(it_acc)={}'.format(type(it_acc)))
print(list(it_acc))/<code>

Output:

<code>type(it_acc)=<class>
[0, 1, 3, 6, 10, 15, 21, 28, 36, 45]/<class>/<code>

(2). chain()

如果需要把若干個序列融合成一個序列的時候，chain()函數可以解決。

<code>it_chain = it.chain([1,2,3], [4,5])
print('type(it_chain)={}'.format(type(it_chain)))
print(list(it_chain))/<code>

Output:

<code>type(it_chain)=<class>
[1, 2, 3, 4, 5]/<class>/<code>

(3). compress()

想要使用掩膜板從一個序列中提取若干元素形成一個新的序列，compress()函數可以實現。

<code>it_comp = it.compress([1,2,3,4,5], [0,1,1,0,0])
print('type(it_comp)={}'.format(type(it_comp)))
print(list(it_comp))/<code>

Output:

<code>type(it_comp)=<class>
[2, 3]/<class>/<code>

(4). dropwhile()

有一個奇怪的需求，對序列的元素逐個遍歷，遇到第一個不滿足條件的元素時，接受該元素及其以後的所有元素形成一個新的序列。這個拗口的描述就是dropwhile()函數實現的功能。

<code>it_dw = it.dropwhile(lambda x:x<3, [1,2,3,4,5,6])
print('type(it_dw)={}'.format(type(it_dw)))
print(list(it_dw))/<code>

Output:

<code>type(it_dw)=<class>
[3, 4, 5, 6]/<class>/<code>

(5). takewhile()

takewhile()函數提供的功能與dropwhile()剛好相反，即接受第一個不滿足條件的元素之前的所有元素。

<code>it_tw = it.takewhile(lambda x: x<3, [1,2,3,4,5])
print('type(it_tw)={}'.format(type(it_tw)))
print(list(it_tw))/<code>

Output:

<code>type(it_tw)=<class>
[1, 2]/<class>/<code>

(6). filterfalse()

如果需要刪除序列中滿足條件的元素，可以選擇filterfalse()函數。

<code>it_ff = it.filterfalse(lambda x:x<3, [1,2,3,4,5,6])
print('type(it_ff)={}'.format(type(it_ff)))
print(list(it_ff))/<code>

Output:

<code>type(it_ff)=<class>
[3, 4, 5, 6]/<class>/<code>

(7). groupby()

將序列中排放在相鄰位置的相同元素歸為同一組,並以該元素為key，以該組元素為value生成詞典，這是groupby()函數的功能。

<code>it_gb = it.groupby([1,2,2,3,3,3,4,5,5,5])
print('type(it_gb)={}'.format(type(it_gb)))
print([k for k,g in it_gb])
it_gb = it.groupby([1,2,2,3,3,3,4,5,5,5])
print([list(g) for k,g in it_gb])/<code>

Output:

<code>type(it_gb)=<class>
[1, 2, 3, 4, 5]
[[1], [2, 2], [3, 3, 3], [4], [5, 5, 5]]/<class>/<code>

(8). islice()

如果需要對序列進行切片，可以使用islice()函數，雖然比Python自帶的切片複雜一點。

<code># 傳參順序：sequence, start, end, step
it_is = it.islice([1,2,3,4,5], 1, 3, 1)
print('type(it_is)={}'.format(type(it_is)))
print(list(it_is))/<code>

Output:

<code>type(it_is)=<class>
[2, 3]/<class>/<code>

(9). starmap()

如果需要對元素為序列的序列進行處理，starmap()可以提供對應的功能。

<code>it_sm = it.starmap(lambda x,y: x*y, [(1,2),(3,4),(5, 6)])
print('type(it_sm)={}'.format(type(it_sm)))
print(list(it_sm))/<code>

Output:

<code>type(it_sm)=<class>
[2, 12, 30]/<class>/<code>

(10). zip_longest()

將兩個序列逐個元素組合起來進行處理也是比較常用的操作，zip_longest()可以提供更加通用的功能。

<code>it_zl = it.zip_longest([1,2,3,4], [5,6], fillvalue=0)
print('type(it_zl)={}'.format(type(it_zl)))
print(list(it_zl))/<code>

Output:

<code>type(it_zl)=<class>
[(1, 5), (2, 6), (3, 0), (4, 0)]/<class>/<code>

(11). 性能對比

就以zip_longest()函數的為例來進行性能對比，我們需要先實現一下提供相同功能的函數。

<code>def my_zip_longest(seq_1, seq_2, fillvalue):
    length = max(len(seq_1), len(seq_2))
    for i in range(length):
        first = seq_1[i] if i        second = seq_2[i] if i        yield first,second/<code>

然後使用notebook的魔術方法來進行如下評測：

<code>%timeit list(it.zip_longest([1,2,3,4,5,6,7,8,9], [1,2,3,4,5], fillvalue=0))
%timeit list(my_zip_longest([1,2,3,4,5,6,7,8,9], [1,2,3,4,5], fillvalue=0))/<code>

Output:

<code>968 ns ± 5.16 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
3.58 µs ± 37.3 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)/<code>

評測數據告訴我們，itertools提供的zip_longest()比自己動手實現的my_zip_longest()要快3倍左右。

3. 排列組合迭代器

“排列組合迭代器”大概是性價比最高的一組接口，因為這些功能要是自己動手實現，且不說速度慢，代碼不易讀，甚至有些實現起來還不那麼容易。

(1). product()

product()函數可以用於計算兩個序列的笛卡爾積。

<code>it_prod = it.product([1,2,3], [4,5,6])
print('type(it_prod)={}'.format(type(it_prod)))
print(list(it_prod))/<code>

Output:

<code>type(it_prod)=<class>
[(1, 4), (1, 5), (1, 6), (2, 4), (2, 5), (2, 6), (3, 4), (3, 5), (3, 6)]/<class>/<code>

(2). permutations()

如果需要對序列元素進行無重複元素的排列，可以使用permutations()，這裡的無重複可不是針對數值而言，而是針對對象的。所以如果想要數值不重複，需要保證輸入序列數值不重複。

<code>it_perm = it.permutations([1,2,3], 2)
print('type(it_perm)={}'.format(type(it_perm)))
print(list(it_perm))/<code>

Output:

<code>type(it_perm)=<class>
[(1, 2), (1, 3), (2, 1), (2, 3), (3, 1), (3, 2)]/<class>/<code>

(3). combinations()

有排列就有組合，combinations()函數就是用來提供組合功能的。

<code>it_comb = it.combinations([1,2,3], 2)
print('type(it_comb)={}'.format(type(it_comb)))
print(list(it_comb))/<code>

Output:

<code>type(it_comb)=<class>
[(1, 2), (1, 3), (2, 3)]/<class>/<code>

(4). combinations_with_replacement()

如果想要進行有重複元素的組合，該怎麼辦？使用combinations_with_replacement()解決。

<code>it_cwr = it.combinations_with_replacement([1,2,3], 2)
print('type(it_cwr)={}'.format(type(it_cwr)))
print(list(it_cwr))/<code>

Output:

<code>type(it_cwr)=<class>
[(1, 1), (1, 2), (1, 3), (2, 2), (2, 3), (3, 3)]/<class>/<code>

(5). 性能對比

組合功能實現起來稍微複雜一點，我們就以組合功能來對比一下itertools提供的工具與手動實現的接口。首先實現一個與combinations()函數功能大體一致的函數。

<code>def my_combinations(seq, times):
    length = len(seq)
    if length < times:
        return
    idx = list(range(times))
    yield tuple(seq[i] for i in idx)
    
    def find_idx_i(idx, times, length):
        for i in range(times-1, -1, -1):
            if idx[i] != i + length - times:
                return i
        return None
    
    while True:
        i = find_idx_i(idx, times, length)
        if i is not None:
            idx[i] += 1
            for j in range(i+1, times):
                idx[j] = idx[j-1]+1
            yield tuple(seq[k] for k in idx) 

        else:
            return
print(list(my_combinations([1,2,3, 4], 3)))/<code>

Output:

<code>[(1, 2, 3), (1, 2, 4), (1, 3, 4), (2, 3, 4)]/<code>

從輸出結果來看，實現的功能正常。

接下來使用notebook提供的魔術方法測試一下運行時間。

<code>%timeit list(it.combinations([1,2,3,4], 3))
%timeit list(my_combinations([1,2,3,4], 3))/<code>

Output:

<code>576 ns ± 12.7 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
7.52 µs ± 96.8 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)/<code>

從對比評測的數據來看，itertools提供的combinations()比自行實現的my_combinations()要快13倍左右。或許是我的實現方案有問題，大家也可以自己實現一個試試看。

經過三個性能評測的對比，很明顯地，itertools提供的接口比自己實現的要快很多，而且調用簡單，代碼易讀。因此，學會itertools工具優化程序中的循環結構，大概可以拯救那些低效緩慢的Python程序。

閱讀更多 青衣極客 的文章

關鍵字: Python 鏡音雙子 低效