Оглавление

Почему я затеял писать про top-K из N? Это довольно смешная история, но давайте сперва решим задачку.

Нахождение top-K из N

Пару слов про кучу

Что за структура данных куча, можно почитать во множестве статей. Вот хорошая, на wiki тоже можно почитать.

Коротко: это бинарное дерево (у каждого узла максимум два потомка), которое отвечает простому условию: все потомки больше родителя (это для min-heap, мы будем использовать именно её). Кучу удобно хранить в одномерном массиве.

Нам понадобятся только два метода работы с кучей: (1) добавление и (2) замена минимального на новый. Я постарался всё прокомментировать в коде, так что, давайте ближе к делу.

Алгоритм нахождения top-K из N с помощью кучи

Я сделал четыре реализации и функцию для тестирования, которая сравнивает результаты работы всех четырёх функций.

#!/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 и только одну операцию…) Тогда я решил всё же написать хоть раз в жизни этот алгоритм руками. Что я и сделал.

Надо сказать, если вы это делаете первый раз, даже хорошо представляя алгоритм, то вряд ли уложитесь в полчаса. Для задачки на собеседовании, — это абсолютно неприемлемая скорость. Если вы планируете идти на интервью в хорошую контору, то обязательно поупражняйтесь с подобными алгоритмами.