Как улучшить этот код для поиска k самых больших элементов массива? - PullRequest
Новая
       50

Как улучшить этот код для поиска k самых больших элементов массива?

0 голосов
/ 11 июля 2020

Следующий код для поиска k самых больших элементов массива вызывает ошибку TLE. Как оптимизировать его, чтобы он работал быстрее? 

import heapq    
for _ in range(int(input())):         
    n,k=map(int,input().split())   
    lists=list(map(int,input().split()))  
    
    heapq.heapify(lists)       
    
    for i in range(k+1):
        klargest=heapq.nlargest(i,lists)  
    
    print(*klargest)

1 Ответ

1 голос
/ 15 июля 2020 
for i in range(k+1):
   klargest=heapq.nlargest(i,lists)

Временная сложность каждой наибольшей операции составляет O (k * log n)), где n - количество элементов в куче. В приведенном выше фрагменте кода эта операция выполняется k + 1 раз для значений [0, k].

Расчет времени для l oop:

итерации значение (Время)

i == 0 (0 * log (n))

i == 1 (1 * log (n))

i == 2 (2 * журнал (n))

....

i == k-1 ((k-1) * log (n))

i == k ((k) * log (n))

Общее время будет суммой времени, затраченного на каждую операцию = (0.log (n)) + (1 * log (n) ) + .... + ((k-1) * log (n)) + ((k) * log (n))

Общее время = (0 + 1 + 2 ... + ( k-1) + k) журнал (n) = ((k (k + 1)) / 2) * журнал (n)

Общее время ~~ O (k ^ 2 * (log (n)))

Вот почему приведенный выше код приводит к TLE.

ОПТИМИЗИРОВАННЫЙ ПОДХОД: 

import heapq    
for _ in range(int(input())):         
    n,k=map(int,input().split())   
    lists=list(map(int,input().split()))  
    
    heapq.heapify(lists)       
    
    for i in range(n-k):
        heapq.heappop(lists)
    klargest = list(lists) # converting heap to list
    print(*klargest)

Поскольку встроенная куча в python является минимальной кучей. Таким образом, приведенный выше код выводит минимум nk элементов из списков. На вывод каждой операции потребуется log (n) раз. Таким образом, общее время будет ~~ (nk) * logn. Оставшиеся k элементов в куче - это самые большие элементы, которые мы хотим найти.

Таким образом, временная сложность для приведенного выше решения составляет O ((nk) * log (n)) == O (nlog (n)) , что является оптимизированной временной сложностью.

...