[转载]如何使用python在2M内存中排序一百万个32位整数

有人开玩笑地问我 如何使用python在2M内存中排序一百万个32位整数.为了应付这个挑战,我学习了一下缓冲I/O.
很 明显,这是一个开玩笑的问题.假设是二进制编码,单单是数据就已经占了4M!唯一的解释就是: 给定一个包含一百万个32位整数的文件,你如何使用最少内存去排序好它们?这可能需要使用某种合并排序的方式,把数据分块在内存排序,并保存到临时文件中 去,最后把临时文件合并获得最终结果.
下面是我的解决方案,稍候会解释.
注意:所有的例子都是使用python3.0. 这个例子的不同之处就是使用file.buffer访问二进制流的方式去访问字符流文件.

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

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

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

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)

在 我的Goole桌面(一个使用了3年,跑着Googlified Linux的个人电脑,用Python3.0的pystones测试得分是34000),在输入文件包含一百万个32位随机整数的情况下,这个程序大概花 了5.4秒.还不太差,如果直接在内存中排序大概需要2秒:

#!/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来说差别并不是很大!