Heaps in practice

Use heapq to find the k smallest numbers and build a priority queue for tasks with (priority, name) tuples.

Two patterns you will reach for often: the top-K heap and the task priority queue. Read through the code, adjust the inputs, and re-run to build intuition.

find_k_smallest and a task priority queue

Python — editable, runs in your browser

import heapq

# ── find_k_smallest ───────────────────────────────────────────────────────────

def find_k_smallest(nums, k):
  """Return the k smallest numbers from nums, sorted ascending.

Uses a max-heap of size k (implemented by negating values).
  Time: O(n log k), Space: O(k).
  """
  if k <= 0:
      return []
  # Build an initial max-heap from the first k elements.
  # We negate values because heapq is a min-heap; negating turns it into
  # a max-heap so heap[0] holds the LARGEST of the k candidates.
  heap = [-x for x in nums[:k]]
  heapq.heapify(heap)

for x in nums[k:]:
      if -x > heap[0]:   # x is smaller than the current largest candidate
          heapq.heapreplace(heap, -x)

return sorted(-v for v in heap)

# ── Priority queue for tasks ──────────────────────────────────────────────────

class TaskQueue:
  """Min-heap priority queue: lower number = higher priority."""

def __init__(self):
      self._heap = []
      self._counter = 0   # tie-break on equal priority (insertion order)

def push(self, priority, name):
      heapq.heappush(self._heap, (priority, self._counter, name))
      self._counter += 1

def pop(self):
      priority, _, name = heapq.heappop(self._heap)
      return priority, name

def peek(self):
      priority, _, name = self._heap[0]
      return priority, name

def __len__(self):
      return len(self._heap)

# ── demos ─────────────────────────────────────────────────────────────────────

print("=== find_k_smallest ===")
data = [7, 2, 9, 1, 5, 3, 8, 4, 6]
for k in [1, 3, 5]:
  result = find_k_smallest(data, k)
  print(f"  k={k}: {result}")

print()
print("=== TaskQueue (lower priority number = more urgent) ===")
tq = TaskQueue()
tq.push(3, "send weekly report")
tq.push(1, "fix production bug")
tq.push(2, "review pull request")
tq.push(1, "respond to incident alert")

while tq:
  pri, name = tq.pop()
  print(f"  [{pri}] {name}")

Why negate for a max-heap

heapq only provides a min-heap. To track the largest of k candidates (so we can replace it when a smaller number comes along), we need a max-heap. The trick: store -x instead of x. The smallest negative is the largest original value. Negate on push, negate on pop.

heapreplace(heap, item) is a combined pop-then-push that is slightly faster than calling them separately — use it inside the scan loop.

The tie-breaking counter

If two tasks have the same priority, comparing (1, "fix bug") and (1, "respond to alert") would try to compare the strings, which works for strings but breaks for arbitrary objects. The counter inserted as the second tuple element guarantees a unique, ordered tie-break and avoids any comparison on the task name itself.

For Dijkstra's algorithm, replace the task name with a (distance, node) tuple. The heap always gives you the unvisited node with the smallest known distance — the same pattern as the priority queue here, just with graph nodes instead of task names.

Where to go next

Next: lab — structure problems — four problems where you choose the right data structure: a min-stack, level-order traversal, k most frequent, and first unique character.

Finished reading? Mark it complete to track your progress.

find_k_smallest and a task priority queue

Why negate for a max-heap

The tie-breaking counter

Where to go next

On this page