Python 节省内存的循环写法 (一)

0 前言

说到处理循环,我们习惯使用 for, while 等,比如依次打印每个列表中的字符:

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

输出:

I
love
python

在打印内容字节数较小时,全部载入内存后,再打印,没有问题。可是,如果现在有成千上百万条车辆行驶轨迹,叫你分析出其中每个客户的出行规律,堵车情况等,假如是在单机上处理这件事。
你可能首先要面临,也可能被你忽视,最后代码都写好后,才可能暴露出的一个问题:outofmemory, 这在实际项目中经常遇到。
这个问题提醒我们,处理数据时,如何写出高效利用内存的程序,就显得很重要。今天,我们就来探讨如何高效利用内存,节省内存同时还能把事情办好。

其实,Python 已经准备好一个模块专门用来处理这件事,它就是 itertools 模块,这里面几个函数的功能其实很好理解。
我不打算笼统的介绍它们所能实现的功能,而是想分析这些功能背后的实现代码,它们如何做到高效节省内存的,Python 内核的贡献者们又是如何写出一手漂亮的代码的,这很有趣,不是吗?

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(dataselectors)

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(predicateiterable)
应用例子:

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(predicateiterable)

应用例子:

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(predicateiterable)

应用例子:

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