Solución:
heapq
es un montón binario, con O (log n) push
y O (log n) pop
. Consulte el código fuente de heapq.
El algoritmo que muestra toma O (n log n) para empujar todos los elementos al montón, y luego O ((nk) log n) para encontrar el k-ésimo elemento más grande. Entonces la complejidad sería O (n log n). También requiere O (n) espacio extra.
Puede hacer esto en O (n log k), usando O (k) espacio extra modificando ligeramente el algoritmo. No soy un programador de Python, así que tendrás que traducir el pseudocódigo:
# create a new min-heap
# push the first k nums onto the heap
for the rest of the nums:
if num > heap.peek()
heap.pop()
heap.push(num)
# at this point, the k largest items are on the heap.
# The kth largest is the root:
return heap.pop()
La clave aquí es que el montón contiene solo los elementos más grandes vistos hasta ahora. Si un elemento es más pequeño que el k-ésimo más grande visto hasta ahora, nunca se coloca en el montón. El peor de los casos es O (n log k).
Realmente, heapq
tiene un heapreplace
método, por lo que podría reemplazar esto:
if num > heap.peek()
heap.pop()
heap.push(num)
con
if num > heap.peek()
heap.replace(num)
Además, una alternativa a empujar la primera k
elementos es crear una lista de los primeros k
artículos y llamar heapify
. Un algoritmo más optimizado (pero aún O (n log k)) es:
# create array of first `k` items
heap = heapify(array)
for remaining nums
if (num > heap.peek())
heap.replace(num)
return heap.pop()
También puedes llamar heapify
en toda la matriz, luego haga estallar la primera n-k
elementos, y luego tome la parte superior:
heapify(nums)
for i = 0 to n-k
heapq.heappop(nums)
return heapq.heappop(nums)
Eso es más sencillo. No estoy seguro de si es más rápido que mi sugerencia anterior, pero modifica la matriz original. La complejidad es O (n) para construir el montón, luego O ((nk) log n) para los pops. Entonces es O ((nk) log n). Peor caso O (n log n).