Heap and Top K Elements in Python: heapq Patterns for Coding Interviews

When a problem asks for the k largest, the k closest or the median so far, a heap is usually the intended tool. A heap hands you the smallest (or largest) item instantly and lets you insert or remove in O(log n), so you never pay to sort everything. This guide explains how Python’s heapq behaves, the three heap patterns that cover most interview questions, and six problems whose code was run and checked against the sample inputs.

How heapq works

A heap is a binary tree stored inside a plain list, arranged so that every parent is smaller than or equal to its children. Python’s heapq module gives you a min-heap: heap[0] is always the smallest element. The rest of the list is only partly ordered, so printing a heap does not show a sorted sequence.

Max-heaps and tuples

heapq has no max-heap, so the usual trick is to store negated values: push -x and negate again when you pop. To attach extra data, push a tuple such as (priority, item). Python compares tuples element by element, which matters in the merge problem below.

Pattern 1: a heap of size k

To track the k largest values you keep a min-heap that never grows beyond k items. Each new value goes in, and if the heap is too big you pop the smallest. What remains are the k largest values seen so far, and the root is the smallest of them, which is exactly the kth largest. Memory stays at O(k) no matter how long the input is.

Kth Largest Element in a Stream (LeetCode 703)

import heapq

class KthLargest:
    def __init__(self, k, nums):
        self.k = k
        self.heap = []
        for n in nums:
            self.add(n)

    def add(self, val):
        heapq.heappush(self.heap, val)
        if len(self.heap) > self.k:
            heapq.heappop(self.heap)
        return self.heap[0]

With k = 3 and starting numbers [4, 5, 8, 2], adding 3, 5, 10, 9 and 4 returns 4, 5, 5, 8 and 8. Each call costs O(log k). The same idea solves Kth Largest Element in an Array (LeetCode 215) in O(n log k); quickselect can do it in O(n) on average but is harder to write correctly under pressure.

Top K Frequent Elements (LeetCode 347)

import heapq
from collections import Counter

def top_k_frequent(nums, k):
    counts = Counter(nums)
    return [n for n, _ in heapq.nlargest(k, counts.items(), key=lambda x: x[1])]

Counting takes O(n). Selecting the k most common among m distinct values takes O(m log k). For [1, 1, 1, 2, 2, 3] and k = 2 the answer is [1, 2]. If the interviewer asks for better than O(n log k), mention bucket sort: put each number in a bucket indexed by its count and read the buckets from the highest count down, which is O(n).

K Closest Points to Origin (LeetCode 973)

import heapq

def k_closest(points, k):
    heap = []
    for x, y in points:
        dist = x * x + y * y          # no square root needed to compare
        heapq.heappush(heap, (-dist, x, y))
        if len(heap) > k:
            heapq.heappop(heap)       # drops the farthest point
    return [[x, y] for _, x, y in heap]

Here we want the k smallest distances, so the heap must be a max-heap on distance, which is why the distance is negated. The farthest of the current k sits at the root and is discarded when a closer point arrives. Skipping the square root is safe because squaring preserves order for non-negative numbers. For [[1,3],[-2,2],[5,8],[0,1]] and k = 2 the result is [-2, 2] and [0, 1] in some order. Time is O(n log k).

Pattern 2: merging sorted sources

When several sorted lists must be combined, the next output value is always one of the current heads. A heap of those heads gives the smallest in O(log k), where k is the number of lists.

Merge k Sorted Lists (LeetCode 23)

import heapq

class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next

def merge_k_lists(lists):
    heap = []
    for i, node in enumerate(lists):
        if node:
            heapq.heappush(heap, (node.val, i, node))
    dummy = tail = ListNode()
    while heap:
        _, i, node = heapq.heappop(heap)
        tail.next = node
        tail = node
        if node.next:
            heapq.heappush(heap, (node.next.val, i, node.next))
    return dummy.next

The middle item i is a tie-breaker. If two nodes have equal values, Python moves on to compare the next tuple element, and two ListNode objects cannot be compared, which raises a TypeError. The unique index prevents that. Merging [1,4,5], [1,3,4] and [2,6] gives [1,1,2,3,4,4,5,6]. With N nodes in total, time is O(N log k) and the heap holds at most k nodes. For the basics of node manipulation see our linked list guide.

Pattern 3: two heaps for a running median

To find the median of a growing stream, split the numbers into a lower half and an upper half. Keep the lower half in a max-heap and the upper half in a min-heap. The median is then at the tops of the heaps.

Find Median from Data Stream (LeetCode 295)

import heapq

class MedianFinder:
    def __init__(self):
        self.small = []   # max-heap (stored negated): lower half
        self.large = []   # min-heap: upper half

    def add_num(self, num):
        heapq.heappush(self.small, -num)
        heapq.heappush(self.large, -heapq.heappop(self.small))
        if len(self.large) > len(self.small):
            heapq.heappush(self.small, -heapq.heappop(self.large))

    def find_median(self):
        if len(self.small) > len(self.large):
            return float(-self.small[0])
        return (-self.small[0] + self.large[0]) / 2

Each new number is pushed through the lower heap into the upper heap, which guarantees every value in small is no larger than every value in large. The size check then keeps small equal in size to large, or one larger. After adding 1, 2 and 3 the medians are 1.0, 1.5 and 2.0. Adding is O(log n) and reading the median is O(1).

Complexity summary

If you need a refresher on reading these notations, start with Big O notation explained.

Common mistakes

• Assuming heapq is a max-heap. It is a min-heap; negate values for the max-heap behaviour.
• Reading heap[1] as the second smallest. Only heap[0] is guaranteed; the rest is partly ordered.
• Using a max-heap of size k when you need the k largest. You need a min-heap, so the smallest of the kept values is the one you evict.
• Pushing objects that cannot be compared. Add a unique tie-breaker such as an index to the tuple.
• Calling sorted() on everything when only k results are needed. That costs O(n log n) instead of O(n log k).

Quick answers

When in doubt

If the question contains the words top, kth, closest or median together with a large or streaming input, think heap first. If the data is already sorted, two pointers or binary search are often simpler. Keep the Python interview cheat sheet nearby for the heapq and collections details.

Also see: Dijkstra’s algorithm and interval problems for related patterns.