Python Program to Find the K Largest Elements Using a Heap
Intermediate Data structures Time O(n log k) Space O(k)
Finds the k largest items in a list with heapq, keeping a min-heap of the k biggest seen so far, so the list never has to be sorted.
To find the k largest items in a list, you could sort the whole list and take the first k. A heap does less work: it keeps only the k biggest items seen so far, and always knows which of them is the smallest.
Python's heapq module keeps an ordinary list arranged as a min-heap: the smallest item is always at index 0, and
adding or removing an item takes about log₂ k steps instead of a full sort. For the top 3 of nine scores, the heap
never holds more than 3. A new score that beats the smallest of those 3 takes its place; any other is skipped.
Program
import heapq
def top_k(items, k):
"""The k largest items, largest first."""
heap = [] # min-heap: heap[0] is the smallest kept
for item in items:
if len(heap) < k:
heapq.heappush(heap, item)
elif item > heap[0]:
# Swap the smallest kept for this bigger item
heapq.heapreplace(heap, item)
return sorted(heap, reverse=True)
scores = [42, 17, 93, 8, 61, 75, 29, 88, 54]
print("Top 3 scores:", top_k(scores, 3))
Open in the editor Download top_k_with_heap.py
Output
Top 3 scores: [93, 88, 75]
How it works
Step through the program line by line: each step below shows up as its line runs.
-
heapqis the standard library's heap module. It works on a plain list. -
heapstarts empty.heapqkeeps it as a min-heap, soheap[0]is always the smallest item in it. The rest isn't fully sorted, only arranged so that finding the next smallest is quick. - The first k items go straight in with
heappush(). - After that, an item bigger than the smallest one kept belongs in the top k, and that
smallest one doesn't any more.
heapreplace()removesheap[0]and adds the new item in one step. An item that isn't bigger can't be in the top k, so it's skipped. - The heap now holds the k largest items.
sorted(heap, reverse=True)lists them largest first. - The program finds the top 3 of nine scores.
| item | What happens | heap afterwards |
|---|---|---|
| 42 | Fewer than 3 kept: push | [42] |
| 17 | Fewer than 3 kept: push | [17, 42] |
| 93 | Fewer than 3 kept: push | [17, 42, 93] |
| 8 | 8 < 17: skip | [17, 42, 93] |
| 61 | 61 > 17: replace 17 | [42, 61, 93] |
| 75 | 75 > 42: replace 42 | [61, 75, 93] |
| 29 | 29 < 61: skip | [61, 75, 93] |
| 88 | 88 > 61: replace 61 | [75, 88, 93] |
| 54 | 54 < 75: skip | [75, 88, 93] |
Here the heap's list happens to stay sorted. With more items it usually isn't, but heap[0] is always the smallest.
Time and space complexity
| Time | O(n log k) |
|---|---|
| Extra space | O(k) |
Each of the n items is compared with heap[0], which takes O(1). Pushing or replacing rearranges a heap of at most
k items, which takes O(log k). So the loop is O(n log k), and sorting the k results at the end adds O(k log k), which
is no more than that. The heap holds k items: O(k) extra space.
Sorting the whole list instead takes O(n log n) time and O(n) space. The gap grows with the list: for the top 10 of a million items, log₂ k is about 3.3 against about 20 for log₂ n, and the heap holds 10 items instead of a sorted copy of a million.
Other ways to write it
With heapq.nlargest()
heapq has this built in. nlargest() works the same way, with a heap of k items, and takes a key argument as
sorted() does.
import heapq scores = [42, 17, 93, 8, 61, 75, 29, 88, 54] print(heapq.nlargest(3, scores))
By sorting
The shortest way, and fine for a small list, but it sorts every item: O(n log n).
scores = [42, 17, 93, 8, 61, 75, 29, 88, 54] print(sorted(scores, reverse=True)[:3])
Try it yourself
- Try k = 1, k = 9 (every score) and k = 12, more than there are.
- Keep
(score, name)tuples, such as(93, "Asha"), to get the names of the top three: tuples compare by their first item first. - Find the 3 lowest scores.
heapqhas only min-heaps, so store-item: thenheap[0]is minus the largest score kept, the one to replace. - Find the 3 longest words in a sentence with
heapq.nlargest(3, words, key=len).