用迭代器模块为Python提速

Python官方文档用"高效的循环"来形容itertools模块，有些tools会带来性能提升，而另外一些tools并不快，只是会节省一些开发时间而已，如果滥用还会导致代码可读性变差。我们不妨把itertools的兄弟们拉出来溜溜。

1. 数列累加

给定一个列表An，返回数列累加和Sn。举例说明：

输入: [1, 2, 3, 4, 5]

返回: [1, 3, 6, 10, 15]

使用accumulate，性能提升了2.5倍

from itertools import accumulate

def _accumulate_list(arr):
    tot = 0
    for x in arr:
        tot += x
        yield tot

def accumulate_list(arr):
    return list(_accumulate_list(arr))

def fast_accumulate_list(arr):
    return list(accumulate(arr))

arr = list(range(1000))

%timeit accumulate_list(arr)
61 µs ± 2.91 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit fast_accumulate_list(arr)
21.3 µs ± 811 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

2. 选择数据

给定一个列表data，一个用0/1表示的列表selectors，返回被选择的数据。举例说明：

输入: [1, 2, 3, 4, 5], [0, 1, 0, 1, 0]

返回: [2, 4]

使用compress，性能提升了2.8倍

from itertools import compress
from random import randint

def select_data(data, selectors):
    return [x for x, y in zip(data, selectors) if y]

def fast_select_data(data, selectors):
    return list(compress(data, selectors))

data = list(range(10000))
selectors = [randint(0, 1) for _ in range(10000)]

%timeit select_data(data, selectors)
341 µs ± 17.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit fast_select_data(data, selectors)
130 µs ± 3.19 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

3. 组合

给定一个列表arr和一个数字k，返回从arr中选择k个元素的所有情况。举例说明：

输入: [1, 2, 3], 2

返回: [(1, 2), (1, 3), (2, 1), (2, 3), (3, 1), (3, 2)]

使用permutations，性能提升了10倍

from itertools import permutations

def _get_permutations(arr, k, i):
    if i == k:
        return [arr[:k]]
    res = []
    for j in range(i, len(arr)):
        arr_cpy = arr.copy()
        arr_cpy[i], arr_cpy[j] = arr_cpy[j], arr_cpy[i]
        res += _get_permutations(arr_cpy, k, i + 1)
    return res

def get_permutations(arr, k):
    return _get_permutations(arr, k, 0)


def fast_get_permutations(arr, k):
    return list(permutations(arr, k))

arr = list(range(10))
k = 5

%timeit -n 1 get_permutations(arr, k)
15.5 ms ± 1.96 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit -n 1 fast_get_permutations(arr, k)
1.56 ms ± 284 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

4. 筛选数据

给定一个列表arr，筛选出所有的偶数。举例说明：

输入: [3, 1, 4, 5, 9, 2]

返回: [(4, 2]

使用filterfalse，性能反而会变慢，所以不要迷信itertools。

from itertools import filterfalse

def get_even_nums(arr):
    return [x for x in arr if x % 2 == 0]

def fast_get_even_nums(arr):
    return list(filterfalse(lambda x: x % 2, arr))

arr = list(range(10000))

%timeit get_even_nums(arr)
417 µs ± 18.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit fast_get_even_nums(arr)
823 µs ± 22.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

5. 条件终止

给定一个列表arr，依次对列表的所有数字进行求和，若遇到某个元素大于target之后则终止求和，返回这个和。举例说明：

输入: [1, 2, 3, 4, 5], 3

返回: 6 (4 > 3，终止)

使用takewhile，性能反而会变慢，所以不要迷信itertools。

from itertools import takewhile

def cond_sum(arr, target):
    res = 0
    for x in arr:
        if x > target:
            break
        res += x
    return res

def fast_cond_sum(arr, target):
    return sum(takewhile(lambda x: x <= target, arr))

arr = list(range(10000))
target = 5000

%timeit cond_sum(arr, target)
245 µs ± 11.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit fast_cond_sum(arr, target)
404 µs ± 13.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

6. 循环嵌套

给定列表arr1，arr2，返回两个列表的所有元素两两相加的和。举例说明：

输入: [1, 2], [4, 5]

返回: [1 + 4， 1 + 5， 2 + 4， 2 + 5]

使用product，性能提升了1.25倍。

from itertools import product

def _cross_sum(arr1, arr2):
    for x in arr1:
        for y in arr2:
            yield x + y

def cross_sum(arr1, arr2):
    return list(_cross_sum(arr1, arr2))


def fast_cross_sum(arr1, arr2):
    return [x + y for x, y in product(arr1, arr2)]

arr1 = list(range(100))
arr2 = list(range(100))

%timeit cross_sum(arr1, arr2)
484 µs ± 16.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit fast_cross_sum(arr1, arr2)
373 µs ± 11.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

7. 二维列表转一维列表

给定二维列表arr，转为一维列表 举例说明：

输入: [[1, 2], [3, 4]]

返回: [1, 2, 3, 4]

使用chain，性能提升了6倍。

from itertools import chain

def _flatten(arr2d):
    for arr in arr2d:
        for x in arr:
            yield x

def flatten(arr2d):
    return list(_flatten(arr2d))


def fast_flatten(arr2d):
    return list(chain(*arr2d))

arr2d = [[x + y * 100 for x in range(100)] for y in range(100)]

%timeit flatten(arr2d)
379 µs ± 15.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit fast_flatten(arr2d)
66.9 µs ± 3.43 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

作者：李小文，先后从事过数据分析、数据挖掘工作，主要开发语言是Python，现任一家小型互联网公司的算法工程师。

Github: https://github.com/tushushu