使用Python在2M内存中排序一百万个32位整数

核心提示： #!/usr/bin/env python3.0import sys, arraya = array.array('i', sys.stdin.buffer.read())a = list(a)a.sort()a = array.array('i', a)a

#!/usr/bin/env python3.0
import sys, array
a = array.array('i', sys.stdin.buffer.read())
a = list(a)
a.sort()
a = array.array('i', a)
a.tofile(sys.stdout.buffer)

回到合并排序方案.头三行非常明显:

#!/usr/bin/env python3.0
import sys, array, tempfile, heapq
assert array.array('i').itemsize == 4

第一行表示我们使用Python3.0.第二行引入将使用的模块.第三行检查到64位系统时中断程序,因为64位系统中'i'的类型值不是表示32位整数;这里我没有考虑代码的移植性.

然后我们定义一个辅助函数,它是一个从文件读入的整数的生成器:

def intsfromfile(f):
　while True:
　　　a = array.array('i')
　　　a.fromstring(f.read(4000))
　　　if not a:
　　　　　break
　　　for x in a:
　　　　　yield x

这 个地方已经调优了算法性能: 它每次读入1000个整数,然后一个个yield.我原来没有用到缓冲 -- 它就每次只从文件读入4个字节,转换成整数,然后yield.但程序跑起来就慢了4倍!要注意的是我们不能使用a.fromfile(f,1000)因为 fromfile()方法在文件没有足够的值的情况下会出错,并且我想要代码能自动适应文件中的整数.(我原先设想大概会写入10,000个整数到一个临 时文件中.)

接着我们进入循环.反复地从输入文件读入10,000个整数,在内存中排序,然后写入临时文件.我们为临时文件增加一个迭代器,通过上面的intsfromfile()函数,使其成为迭代器中的一个列表,以备在随后的合并阶段使用.

iters = []
while True:
　a = array.array('i')
　a.fromstring(sys.stdin.buffer.read(40000))
　if not a:
　　　break
　f = tempfile.TemporaryFile()
　array.array('i', sorted(a)).tofile(f)
　f.seek(0)
　iters.append(intsfromfile(f))

要注意的是,为了应对包含一百万个值的输入,程序将会创建100个包含10,000个值的临时文件.

最后我们合并这些文件(每个都已经排序). heapq模块有一个非常实用的函数: heapq.merge(iter1, iter2, ...)它把多个已排序的输入值(如本例中的)合并排序，返回一个有序的迭代器。

a = array.array('i')
for x in heapq.merge(*iters):
　a.append(x)
　if len(a) >= 1000:
　　　a.tofile(sys.stdout.buffer)
　　　del a[:]
if a:
　a.tofile(sys.stdout.buffer)

这里也是一个必不可少使用缓冲I/O的地方:当一个值可用的时候就写入文件,这样会成倍地减慢算法速度.因此,通过简单地增加输入输出缓冲,可以获得上十倍的速度提升!

这就是程序的主要思路了.

另外我们还能体验到heapq模块的强大，它包含了在输出阶段需要的合并迭代器功能。当然，array模块管理二进制数据的便利也是令人印象深刻的。

最后,通过这个例子让你们知道,Python 3.0相对于Python2.5来说差别并不是很大!