Сортировка миллиона 32-битных int'ов в 2 мегабайтах памяти на Питоне

Мой перевод статьи Гвидо ван Россума:

Меня тут в шутку спросили: смогу ли я отсортировать миллион 32-битных int'ов в 2 мегабайтах памяти на Питоне. Во время размышления, мне пришло в голову задействовать механизм ввода-вывода с использованием буферной памяти.

Вообще, это именно шуточный вопрос — одни только данные займут 4 мегабайта, при условии бинарного представления! Правда, можно пойти на хитрость — взять файл, содержащий миллион 32-битных int'ов. Как же отсортировать их, используя минимальное количество памяти? Это должна быть какая-то разновидность сортировки слиянием, в которой небольшие куски данных сортируются и записываются во временный файл, после чего происходит слияние временных файлов для получения окончательного результата.

Вот мое решение:

Copy Source | Copy HTML

1. #!/usr/bin/env python3.0
2. import sys, array, tempfile, heapq
3. assert array.array('i').itemsize == 4
4.  
5. def intsfromfile(f):
6.  while True:
7.     a = array.array('i')
8.     a.fromstring(f.read(4000))
9.     if not a:
10.         break
11.     for x in a:
12.         yield x
13.  
14. iters = []
15. while True:
16.  a = array.array('i')
17.  a.fromstring(sys.stdin.buffer.read(40000))
18.  if not a:
19.      break
20.  f = tempfile.TemporaryFile()
21.  array.array('i', sorted(a)).tofile(f)
22.  f.seek(0)
23.  iters.append(intsfromfile(f))
24.  
25. a = array.array('i')
26. for x in heapq.merge(*iters):
27.  a.append(x)
28.  if len(a) >= 1000:
29.      a.tofile(sys.stdout.buffer)
30.      del a[:]
31. if a:
32.  a.tofile(sys.stdout.buffer)

На моей 3х-летней персоналке под управлением Linux, выполнение заняло 5.4 секунды при чтении данных из входного файла, содержащего ровно миллион случайных 32-битных int'ов. Это не так уж плохо, учитывая, что при чтении данных прямо из памяти сортировка занимает 2 секунды:

Copy Source | Copy HTML

1. #!/usr/bin/env python3.0
2. import sys, array
3. a = array.array('i', sys.stdin.buffer.read())
4. a = list(a)
5. a.sort()
6. a = array.array('i', a)
7. a.tofile(sys.stdout.buffer)

Вернемся к нашей сортировке. Первые 3 строки очевидны:

Copy Source | Copy HTML

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

Первая строка указывает, что мы используем Python 3.0. Вторая строка импортирует необходимые модули. Третья строка вызывает прерывание на 64-битных системах, где int не 32-битный — передо мной не ставилась задача написать портируемый код.

Затем мы объявили функцию, которая считывает int'ы из файла:

Copy Source | Copy HTML

1. def intsfromfile(f):
2.  while True:
3.      a = array.array('i')
4.      a.fromstring(f.read(4000))
5.      if not a:
6.          break
7.      for x in a:
8.          yield x

Это то самое место, где имеет место быть оптимизация алгоритма: он считывает до 1000 int'ов за раз, выдавая их по одному. Вначале я не использовал буферизацию — функция просто считывала 4 байта из файла, преобразовывала в целое и возвращала. Но это работало в 4 раза медленнее! Обратите внимание, что мы не можем использовать a.fromfile(f, 1000), так как метод fromfile() вернет ошибку, если в файле недостаточно данных, а я хотел, чтобы код сам адаптировался к любому количеству int'ов в файле.

Далее следует цикл ввода. Он циклически считывает кусок из 10'000 int'ов из входного файла, сортирует их в памяти и записывает во временный файл. Затем мы добавим итератор для временного файла в список итераторов, который понадобится нам на заключительном этапе, используя описанную выше функцию intsfromfile().

Copy Source | Copy HTML

1. iters = []
2. while True:
3.  a = array.array('i')
4.  a.fromstring(sys.stdin.buffer.read(40000))
5.  if not a:
6.      break
7.  f = tempfile.TemporaryFile()
8.  array.array('i', sorted(a)).tofile(f)
9.  f.seek(0)
10.  iters.append(intsfromfile(f))

Соответственно, для входного миллиона int'ов будет создано 100 временных файлов по 10'000 значений в каждом.

Наконец, мы объединяем все эти файлы (каждый из них уже отсортирован) вместе. Модуль heapq содержит замечательную функцию для решения этой задачи: heapq.merge(iter1, iter2, ... ), возвращающую итератор, который пропускает входные параметры в такой последовательности, при которой каждый входной параметр сам передает свое значение в нужном порядке (как в данном случае).

Copy Source | Copy HTML

1. a = array.array('i')
2. for x in heapq.merge(*iters):
3.  a.append(x)
4.  if len(a) >= 1000:
5.      a.tofile(sys.stdout.buffer)
6.      del a[:]
7. if a:
8.  a.tofile(sys.stdout.buffer)

Это еще один случай, при котором буферизированный ввод\вывод оказывается необходимым: запись каждого отдельного значения в файл, как только оно становится доступным, снижает темп работы алгоритма вдвое. Используя буферизацию на входе и выходе, нам удалось увеличить скорость выполнения в 10 раз!