378 - Kth Smallest Element in a Sorted Matrix

Difficulty: Medium | Pattern: Binary Search / Min-Heap | Company tags: Amazon, Facebook, Twitter

Problem Statement

Given an n x n matrix where each of the rows and columns is sorted in ascending order, return the k-th smallest element in the matrix.

Note that it is the k-th smallest element in the sorted order, not the k-th distinct element.

You must find a solution with a memory complexity better than O(n²).

Example 1:

Input: matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
Output: 13

Example 2:

Input: matrix = [[-5]], k = 1
Output: -5

Approach 1: Min-Heap — O(k log n)

Key insight: Start with all first-column elements in a min-heap. Pop the smallest k times; each time push the next element from that row.

import heapq

def kthSmallest(matrix: list[list[int]], k: int) -> int:
    n = len(matrix)
    heap = [(matrix[r][0], r, 0) for r in range(n)]
    heapq.heapify(heap)
    
    for _ in range(k - 1):
        val, r, c = heapq.heappop(heap)
        if c + 1 < n:
            heapq.heappush(heap, (matrix[r][c+1], r, c+1))
    
    return heapq.heappop(heap)[0]

Approach 2: Binary Search — O(n log(max-min))

Key insight: Binary search on the value range [min, max]. For a given mid, count how many elements are <= mid using the sorted property (start from bottom-left). The k-th smallest is the smallest mid where count >= k.

def kthSmallest(matrix: list[list[int]], k: int) -> int:
    n = len(matrix)
    
    def count_lte(target):
        count = 0
        r, c = n - 1, 0
        while r >= 0 and c < n:
            if matrix[r][c] <= target:
                count += r + 1
                c += 1
            else:
                r -= 1
        return count
    
    lo, hi = matrix[0][0], matrix[n-1][n-1]
    while lo < hi:
        mid = (lo + hi) // 2
        if count_lte(mid) >= k:
            hi = mid
        else:
            lo = mid + 1
    
    return lo

Algorithm Flow

Dry Run (Min-Heap)

matrix = [[1,5,9],[10,11,13],[12,13,15]], k=8

Initial heap: [(1,0,0),(10,1,0),(12,2,0)]

Pop 7 times, push next in row each time. 8th pop gives 13 ✓

Complexity

Approach	Time	Space
Min-Heap	O(k log n)	O(n)
Binary Search	O(n log(max-min))	O(1)

Key Terms

Term	Definition
Min-heap	A priority queue that always pops the smallest element first; here seeded with one element per row.
Binary search on answer	Searching over the range of possible values (not indices) to find the smallest value satisfying a predicate.
Monotonic predicate	A condition (count_lte(mid) >= k) that is false then true as mid increases, enabling binary search.
Sorted matrix (row & column sorted)	A matrix where both rows and columns are individually sorted, allowing staircase-style traversal.

FAQ

Q: Why does binary search work here even though the matrix isn't a single sorted array? A: Binary search operates on the value range [min, max], not on matrix indices. Because count_lte(mid) is monotonic in mid, we can binary search for the smallest value whose count of elements <= mid is at least k.

Q: Why start the counting walk from the bottom-left corner? A: From the bottom-left, moving right increases values and moving up decreases them, so each comparison either adds a whole column's worth of counts or eliminates a row — giving an O(n) count per check.

Q: Which approach is better for interviews, heap or binary search? A: Binary search is generally preferred since it hits O(1) extra space and a better time bound O(n log(max-min)); mention the heap approach first since it's more intuitive, then optimize.

Q: What if k equals n*n (the last element)? A: Both approaches still work — the heap will pop until exhausting all elements, and binary search converges to matrix[n-1][n-1].

Q: Can this be extended to k-th smallest across multiple sorted matrices or lists? A: Yes — the heap approach generalizes directly to k-way merge of sorted lists (same idea as merging k sorted lists).

Quick Revision

• Problem: find k-th smallest value in an n x n matrix sorted by rows and columns.
• Heap approach: seed heap with first column, pop k times, push next element in popped row each time — O(k log n).
• Binary search approach: search value range [matrix[0][0], matrix[n-1][n-1]] using a monotonic count_lte(mid) >= k predicate — O(n log(max-min)), O(1) space.
• Counting walk starts at bottom-left corner to count elements <= mid in O(n).
• Binary search converges when lo == hi, which is guaranteed to be a matrix value.
• Prefer binary search when space is constrained; prefer heap when k is small.
• Do not confuse this with finding the k-th distinct element — duplicates count individually.

Related Problems

• 215 - Kth Largest Element in an Array — same k-th order statistic idea via heap/quickselect.
• Merge k Sorted Lists — shares the heap-based k-way merge pattern.
• Median of Two Sorted Arrays — another binary-search-on-answer problem over sorted structures.