ByteFlow

Сортировка миллиона 32-битных integer в 2МБ памяти используя Python

Кто-то шутя спросил меня, как бы я сортировал миллион 32-битных integer-ов на Python, используя не больше 2МБ памяти? Принимая вызов, я изучил кое-что о буферном вводе-выводе (buffered I/O).

Очевидно, это шуточный вопрос - ведь все данные будут сами занимать до 4 мегабайт в бинарном виде! Но есть еще возможная интерпретация: допустим, исходный файл содержит миллион 32-битных integer-ов. Как бы вы отсортировали их, используя память по-минимуму? Это была бы в некотором роде соединительная сортировка (merge sort), в которой маленькие куски данных сортируются в памяти и записываются во временный файл, а потом временные файлы соединяются в выходной объект.

Вот решение:

Примечание: все примеры используют Python 3.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)

На моем 3-летнем ПК под управлением Linux скрипт с исходным файлом, содержащим ровно 1000000 32-битных произвольных integer, занял примерно 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)

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

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

В первой строке указано использовать Python 3.0. Вторая строка импортирует необходимые модули. В третьей строке указывается размер элемента массива, для 64-битных систем в которых тип ‘i’ не является 32-битным int. Следовательно, тут никаких попыток сделать код портируемым мною не предпринимается.

После этого создаем помощник-генератор который читает integerы из файла и отдает их по-одному:

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 integer-ов сразу, и отдает их один за одним. Сначала я написал это не используя буфер,- алгоритм читал по 4 байта из файла, конвертировал в integer, и выдавал результат. Но это работало примерно в четыре раза медленней! Учтите что мы не можем использовать a.fromfile(f, 1000) потому что метод fromfile() горько жалуется что в файле недостаточно данных, а я хочу чтобы код автоматически подстраивался к любому количеству integer-ов в файле. (Получается, мы пишем примерно 10000 integer-ов в типичный временный файл.)

После этого идет входящий цикл. Он читает куски по 10000 integer-ов со входящего файла, сортирует их в памяти, и записывает их во временный файл. После этого мы добавляем итератор поверх этого временного файла, используя функцию intsfromfile(), чтобы сделать список итераторов которые мы будем использовать на последующем этапе слияния (merge).

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 временных файлов, каждый из которых содержит 10000 чисел.

Наконец, мы соединяем вместе все эти файлы (каждый из которых отсортирован). В модуле 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)

Это еще один пример где буферный ввод-вывод оказался необходим: Запись в файл входящего значения, как только оно появляется, замедляет алгоритм примерно в два раза. Таким образом, просто добавляя буферы ввода и вывода, мы получили десятикратный прирост скорости!

Это и есть мораль всей истории.

Второй урок - хвала heapq модулю, который содержит функцию слияния итераторов, нужную при выдаче данных. Кроме того, давайте не будем забывать о пользе модуля array для управления двоичными данными.

И, в конце концов, этот пример показывает что Python 3.0 не так уж сильно отлчается от Python 2.5 ;)

Оригинал на английском тут.