← Heaps

Drill #60 — Kth smallest in sorted matrix

Medium Heap + BFS Heaps

In plain English: In a grid where rows and columns are sorted, find the kth smallest number without sorting everything.

Start from the top-left corner and use a min-heap to always expand the smallest unvisited cell — like BFS ordered by value.

Prompt

Given an n x n matrix where each row and column is sorted in ascending order, find the kth smallest element.

Try to write it from scratch before scrolling down.

Solution

import heapq

def kth_smallest(matrix, k):
    n = len(matrix)
    # (value, row, col) — start with first element of each row
    heap = [(matrix[i][0], i, 0) for i in range(min(n, k))]
    heapq.heapify(heap)

    val = 0
    for _ in range(k):
        val, r, c = heapq.heappop(heap)
        if c + 1 < n:
            heapq.heappush(heap, (matrix[r][c + 1], r, c + 1))
    return val

# Test: kth_smallest([[1,5,9],[10,11,13],[12,13,15]], 8) == 13
O(k log n) time · O(n) space

Related Micro Drills

← Drill #59 Drill #61 →