~/problems / Heaps / Heaps and priority queues

Basics: the k smallest values with heapq

easy basics ~10 min

Write k_smallest(nums: list[int], k: int) -> list[int] that returns the k smallest values of nums in ascending order, using Python's heapq module.

k_smallest([7, 2, 9, 2, 5], 3)   # [2, 2, 5]   (duplicates count separately)
k_smallest([4, -1, 3], 5)        # [-1, 3, 4]  (k is bigger than the list: return everything)
  • 0 <= k, and k may be larger than len(nums). k == 0 or an empty list gives [].
  • Values may be negative and may repeat.
  • Don't modify nums.
  • The tests use up to a million values with k in the thousands, so scanning for the minimum k times (O(n·k)) is too slow. Aim for O(n + k log n) or O(n log k). (Sorting everything also works; the point here is to practise the heap.)

Two good approaches: heapq.heapify a copy and heappop k times, or keep a max-heap of size k (store negated values, since heapq is a min-heap only). heapq.nsmallest exists too, but write it yourself first.

Show hint

heapq.heapify turns a list into a min-heap in O(n), and each heappop hands you the next-smallest value in O(log n).

Topic: Heaps and priority queues. heapq: top-k, k-way merge, two heaps for a running median.

Read the visual guide
0:00
Ctrl ' run · Ctrl ↵ submit
esc