На этом шаге мы рассмотрим функции, позволяющие решить эту задачу.

Задача

    Вы хотите создать список N максимальных или минимальных элементов коллекции.

Решение

    У модуля heapq есть две функции, nlargest() и nsmallest(), которые делают именно то, что вам нужно. Например:

import heapq
nums = [1, 8, 2, 23, 7, -4, 18, 23, 42, 37, 2] 
print(heapq.nlargest(3, nums)) # Выведет [42, 37, 23] 
print(heapq.nsmallest(3, nums)) # Выведет [-4, 1, 2]

    Обе функции также принимают параметр key, который позволяет использовать их с более сложными структурами данных. Например:

>>> import heapq
>>> portfolio = [
	{'name': 'IBM', 'shares': 100, 'price': 91.1},
	{'name': 'AAPL', 'shares': 50, 'price': 543.22},
	{'name': 'FB', 'shares': 200, 'price': 21.09},
	{'name': 'HPQ', 'shares': 35, 'price': 31.75},
	{'name': 'YHOO', 'shares': 45, 'price': 16.35},
	{'name': 'ACME', 'shares': 75, 'price': 115.65}
	]
>>> cheap = heapq.nsmallest(3, portfolio, key=lambda s: s['price'])
>>> cheap
[{'name': 'YHOO', 'shares': 45, 'price': 16.35}, 
{'name': 'FB', 'shares': 200, 'price': 21.09}, 
{'name': 'HPQ', 'shares': 35, 'price': 31.75}]
>>> expensive = heapq.nlargest(3, portfolio, key=lambda s: s['price'])
>>> expensive
[{'name': 'AAPL', 'shares': 50, 'price': 543.22}, 
{'name': 'ACME', 'shares': 75, 'price': 115.65}, 
{'name': 'IBM', 'shares': 100, 'price': 91.1}]
>>> 

Обсуждение

    Если вы ищете N наименьших или N наибольших элементов, причем N невелико, по сравнению с общим размером коллекции, эти функции покажут великолепную производительность. "Под капотом" они начинают работу с конвертирования данных в список, где данные упорядочены, как в куче. Например:

>>> nums = [1, 8, 2, 23, 7, -4, 18, 23, 42, 37, 2]
>>> import heapq
>>> heap = list(nums)
>>> heapq.heapify(heap)
>>> heap
[-4, 2, 1, 23, 7, 2, 18, 23, 42, 37, 8]
>>> 

    Самая важная возможность кучи состоит в том, что heap[0] всегда будет наименьшим элементом. Кроме того, последующие элементы могут быть легко найдены с помощью метода heapq.heappop(), который удаляет первый элемент и заменяет его следующим наименьшим элементом (это требует O(log N) операций, где N - размер кучи). Например, чтобы найти три наименьших элемента, вы могли бы сделать так:

>>> heapq.heappop(heap)
-4
>>> heapq.heappop(heap)
1
>>> heapq.heappop(heap)
2
>>>

    Функции nlargest() и nsmallest() лучше всего подходят, если вы пытаетесь найти относительно небольшое количество элементов. Если же вы просто хотите найти один наибольший или наименьший элемент (N = 1), функции min() и max() будут быстрее. Похожим образом, если N сопоставимо с размером самой коллекции, обычно будет быстрее отсортировать их и взять срез (то есть выполнить sorted(items)[:N] или sorted(items)[-N:]). Стоит отметить, что реализация nlargest() и nsmallest() в модуле heapq работает гибко и выполняет некоторые из этих оптимизаций самостоятельно (например, использует сортировку, если размер N близок к размеру входных данных).

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

    На следующем шаге мы рассмотрим реализацию очереди с приоритетом.

Предыдущий шаг Содержание Следующий шаг