Как стать автором
Обновить

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

Время на прочтение4 мин
Количество просмотров8.5K
Мой перевод статьи Гвидо ван Россума:

Меня тут в шутку спросили: смогу ли я отсортировать миллион 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 раз!
Теги:
Хабы:
Всего голосов 84: ↑79 и ↓5+74
Комментарии60

Публикации

Истории

Работа

Data Scientist
68 вакансий

Ближайшие события

2 – 18 декабря
Yandex DataLens Festival 2024
МоскваОнлайн
11 – 13 декабря
Международная конференция по AI/ML «AI Journey»
МоскваОнлайн
25 – 26 апреля
IT-конференция Merge Tatarstan 2025
Казань