Почему я затеял писать про top-K из N? Это довольно смешная история, но давайте сперва решим задачку.
Что за структура данных куча, можно почитать во множестве статей. Вот хорошая, на wiki тоже можно почитать.
Коротко: это бинарное дерево (у каждого узла максимум два потомка), которое отвечает простому условию: все потомки больше родителя (это для min-heap, мы будем использовать именно её). Кучу удобно хранить в одномерном массиве.
Нам понадобятся только два метода работы с кучей: (1) добавление и (2) замена минимального на новый. Я постарался всё прокомментировать в коде, так что, давайте ближе к делу.
Я сделал четыре реализации и функцию для тестирования, которая сравнивает результаты работы всех четырёх функций.
#!/usr/bin/python2
# coding: utf8
import random
import heapq
'''
Способ №0: без всякой кучи.
Эту функцию сделаем просто для проверки результатов
'''
def stupid_top(k, it):
return sorted(it, key=lambda x: -x)[:k]
'''
Дальше начинаются методы с использованием кучи
'''
'''
Способ №1: просто используем библиотечную функцию
'''
def straightforward_top(k, it):
return heapq.nlargest(k, it)
'''
Способ №2: используем библиотечные функции для
работы с кучей.
Нам понадобится две функции:
- пока мы не набрали k элементов, мы безусловно
добавляем в кучу новые элементы
- когда куча доросла до k, мы добавляем новый, только
если он меньше минимального в куче
Напомню, что у нас minheap, т.е. в корне именно
минимальный элемент.
'''
def lib_top(k, it):
if k == 0:
return []
heap = []
for i in it:
if len(heap) < k:
heapq.heappush(heap, i)
elif i > heap[0]:
heapq.heappushpop(heap, i)
return heap
'''
Способ №3: тоже самое, что и №2, но сами реализуем работу с кучей,
не используем библиотечные функции.
'''
'''
Памятка: нумерация в кучe
для элемента i:
левый потомок 2 * i + 1
правый потомок 2 * i + 2
родитель (i - 1) // 2
родитель всегда меньше потомков (у нас minheap)
элемент i = 0 — корень
'''
def heappush(h, e):
# добавление нового элемента (куча увеличивается):
# добавляем в конец и пусть всплывает к корню
i = len(h) # индекс нового элемента
h.append(e) # добавляем новый элемент
while i > 0:
p = (i - 1) // 2 # индекс родителя
if h[p] <= h[i]:
break # условие кучи выполнено
else:
(h[p], h[i]) = (h[i], h[p])
i = p
def heappushpop(h, e):
# ставим новый элемент в корень (вместо текущего корня)
# и восстанавливаем кучу:
# меняем с наименьшим из потомков
# если этот потомок меньше нашего элемента
i = 0
h[0] = e
while True:
# индексы потомков
p1 = i * 2 + 1
p2 = p1 + 1
# индекс наименьшего (искомого)
smallest = i
if p1 < len(h):
if h[p1] < h[smallest]:
smallest = p1
if p2 < len(h):
if h[p2] < h[smallest]:
smallest = p2
if smallest == i:
break # перестановка не нужна, условие кучи выполнено
(h[smallest], h[i]) = (h[i], h[smallest])
i = smallest # продолжаем от новой ноды
def top(k, it):
if k == 0:
return []
heap = []
for i in it:
if len(heap) < k:
heappush(heap, i)
elif i > heap[0]:
heappushpop(heap, i)
return heap
'''
Тесты:
10000 раз повторяем тест
- заполняем массив случайной длины случайными элементами
- загадываем случайное k
- прогоняем все четыре алгоритма
- падаем, если результаты отличаются
'''
def runtests():
for i in range(10000):
print "=" * 40
a = map(
lambda x: random.randrange(41),
range(random.randrange(41))
)
k = random.randrange(18)
print "%3d %r" % (k, a)
w = stupid_top(k, a)
x = straightforward_top(k, a)
y = lib_top(k, a)
z = top(k, a)
print w
print x
print y
print z
assert sorted(w) == sorted(x) == sorted(y) == sorted(z)
print "OK"
if __name__ == '__main__':
runtests()
Надеюсь, всё понятно.
Мне много раз задавали эту задачку на собеседованиях. Я всегда отвечал смело: надо использовать кучу. И этот ответ принимался. Причём, я был уверен, что я знаю, как именно её надо использовать.
Я и сам иногда задавал эту задачку кандидатам. И если человек ответствовал: «кучей», то я оставался полностью удовлетворён.
Прошло больше десяти лет, как меня первый раз спросили эту задачку. И вот, на очередном собеседовании ситуация повторяется: мне задают вопрос, я отвечаю «кучей». Но тут мне оппонент говорит: «а напиши код». Я стал писать, и, конечно, стал путаться. Оппонент увидел правильное начало, и, видимо, решив не ждать, когда я всё допишу, рассказал мне всё, чего я пока не написал. Он сказал: «да-да, я понял, сейчас ты сделаешь так-то и так-то и всё будет готово; ответ засчитан».
Но я после этого призадумался крепко. Уже идя с собеседования, я думал над его решением и понимал, что оно не правильное. (Он использовал max-heap и только одну операцию…) Тогда я решил всё же написать хоть раз в жизни этот алгоритм руками. Что я и сделал.
Надо сказать, если вы это делаете первый раз, даже хорошо представляя алгоритм, то вряд ли уложитесь в полчаса. Для задачки на собеседовании, — это абсолютно неприемлемая скорость. Если вы планируете идти на интервью в хорошую контору, то обязательно поупражняйтесь с подобными алгоритмами.