Найти минимум и сделать некоторые вычисления в списке, используя структуру данных кучи

Question

Я изучаю кучу и попробовал вопрос в hackarank.

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

Я закончил кодировать решение. Однако, пожалуйста, укажите любые улучшения, которые я могу сделать, чтобы код работал быстрее.

Мой код:

import heapq as heap

data = map (int, raw_input ().strip ().split ())
N, K = data [0], data [1]

cookies = map (int, raw_input ().strip ().split ()) 
heap.heapify (cookies)
numOps = 0
possibility = False

while cookies [0] < K:
    if N == 1:
        possibility = True
        break
    leastSweetCookies = heap.nsmallest (2, cookies)
    heap.heapreplace (cookies, leastSweetCookies [0] + 2 * leastSweetCookies [1])
    heap.heappop (cookies)
    numOps += 1
    N -= 1
if possibility == False: print numOps
else: print -1

Answer 1

Эти три строки:

leastSweetCookies = heap.nsmallest (2, cookies)
heap.heapreplace (cookies, leastSweetCookies [0] + 2 * leastSweetCookies [1])
heap.heappop (cookies)

эквивалентны:

sw1 = heap.heappop(cookies);
sw2 = heap.heappop(cookies);
heap.heappush(cookies, sw1 + 2*sw2);

В вашем коде вызов nsmallest потенциально перебирает всю кучу, чтобы найти 2 самых маленьких элемента. Затем O (log n), чтобы заменить верхний элемент, и O (log n), чтобы вытолкнуть наименьшее.

В коде замены это O (log n) для каждого из двух всплывающих окон и O (log n) для push. Итак:

ваш код: O (n) + O (журнал n) + O (журнал n)

мой код: O (log n) + O (log n) + O (log n)