Kth Smallest in Sorted Matrix LeetCode 378

LeetCode 378 — Kth Smallest Element in a Sorted Matrix — presents an n x n matrix where each row is sorted in ascending order and each column is sorted in ascending order. Your task is to return the kth smallest element in the matrix (1-indexed). The element is not the kth distinct value — duplicates count separately. You may not simply flatten and sort the matrix, because the follow-up challenge asks you to use the structural property that rows and columns are both sorted.

The two-dimensional sorted structure gives the matrix a unique shape: matrix[0][0] is the global minimum and matrix[n-1][n-1] is the global maximum. Any element is smaller than all elements to its right in the same row and all elements below it in the same column. This cross-row-and-column ordering is stronger than a simple row-sorted matrix and enables efficient counting algorithms that would not work on row-sorted-only arrays.

This problem sits at the intersection of heaps, binary search, and two-dimensional search techniques. Mastering both the min-heap and the binary search approaches builds intuition that transfers directly to problems like Find Median From Data Stream, Search a 2D Matrix, and any problem involving selection-rank queries on structured data.

The min-heap approach treats each row of the matrix as a sorted list and merges them using a priority queue. Initialize the heap with the first element of each row: push (value, row, col) for every row at column 0. Since each row is sorted, the globally smallest element must be among these first-column candidates, and the heap will always return the current minimum.

Pop from the heap k times. On each pop, you extract the current minimum element. If the popped element came from column col in row row, push the next element from that same row — (matrix[row][col+1], row, col+1) — if col+1 is still within bounds. The kth pop yields the kth smallest element. The heap size is at most n (one entry per row) at all times, so each push and pop costs O(log n). Total time for k pops is O(k log n) and space is O(n) for the heap.

This approach is intuitive for anyone familiar with the merge k sorted lists problem. Each row is a sorted list, and the heap efficiently selects the next minimum across all lists without scanning every element. It is especially effective when k is small relative to n*n — for example, finding the minimum (k=1) or a small percentile — because you stop as soon as you have popped k times.

The binary search approach does not search for the position of an element — it searches for the value of the kth smallest element directly. Set lo = matrix[0][0] (the global minimum) and hi = matrix[n-1][n-1] (the global maximum). The answer lies somewhere in this range. For each candidate midpoint mid = (lo + hi) // 2, count the number of elements in the matrix that are less than or equal to mid.

If that count is less than k, the kth smallest element must be larger than mid, so set lo = mid + 1. If the count is greater than or equal to k, the kth smallest element could be mid or something smaller, so set hi = mid. Continue until lo equals hi. At that point, lo is the smallest value for which the count of elements <= lo is at least k — which is exactly the kth smallest element.

The outer binary search runs O(log(max - min)) iterations, where max and min are the largest and smallest matrix values. Each iteration invokes a counting subroutine that must count elements <= mid in the entire matrix. Doing this naively would cost O(n²) per iteration, making the total O(n² log(max-min)) which is no better than sorting. The key is the O(n) staircase counting trick described in the next section.

The staircase counting technique exploits both row sorting and column sorting simultaneously. Start at the bottom-left corner of the matrix: row = n-1, col = 0. At each step, compare matrix[row][col] to mid. If matrix[row][col] <= mid, then every element in the current column from row 0 up through row row is also <= mid — because the column is sorted ascending and all those elements are smaller than matrix[row][col]. Add row+1 to the count and move one step right: col += 1.

If matrix[row][col] > mid, the current element is too large. Move one step up: row -= 1. Continue until col >= n or row < 0. The result is the total count of elements <= mid. Each step either increases col by 1 or decreases row by 1, and col starts at 0 while row starts at n-1, so the total number of steps is at most 2n - 1. The counting subroutine therefore runs in O(n) time.

Combining the O(log(max-min)) outer binary search with the O(n) counting subroutine gives an overall time complexity of O(n log(max-min)) and O(1) extra space. For a 300 x 300 matrix with 32-bit integers, log(max-min) is at most 32, so the algorithm performs at most 300 * 32 = 9600 element comparisons — an extremely efficient bound.

The binary search loop is while lo < hi, not while lo <= hi. This is important because the loop must terminate with lo == hi rather than with lo and hi crossing. Setting lo = mid + 1 when count < k is safe because mid cannot be the answer if fewer than k elements are <= mid. Setting hi = mid when count >= k is safe because hi must remain a valid candidate — we cannot discard mid since it might be exactly the answer.

The standard half-open binary search variant lo = (lo + hi) // 2 computes a mid that is strictly less than hi when lo < hi. This prevents infinite loops: if hi = lo + 1, then mid = lo, and the loop either sets lo = lo + 1 = hi (terminating) or leaves hi = lo (also terminating). Without this property the loop could cycle forever. Using (lo + hi + 1) // 2 instead would break termination in the hi = mid branch.

A subtle invariant: at every iteration, the answer always lies in the closed interval [lo, hi]. When lo == hi, the interval collapses to a single value which must be the answer. The proof relies on the fact that count(matrix[0][0]) >= 1 and count(matrix[n-1][n-1]) == n*n >= k, so the invariant holds at initialization. Each branch preserves it by maintaining lo as a value where count < k is impossible and hi as a value where count >= k holds.

Python binary search solution: def kthSmallest(matrix, k): n = len(matrix); lo, hi = matrix[0][0], matrix[n-1][n-1]; while lo < hi: mid = (lo+hi)//2; count = 0; r, c = n-1, 0; while r >= 0 and c < n: if matrix[r][c] <= mid: count += r+1; c += 1; else: r -= 1; if count < k: lo = mid+1; else: hi = mid; return lo. Time O(n log(max-min)), space O(1). Python heap solution uses heapq.heappush/heappop with (val, row, col) tuples, pushing matrix[i][0] for all rows initially, then popping k times and pushing the next column element each time.

Java binary search solution: public int kthSmallest(int[][] matrix, int k) { int n = matrix.length, lo = matrix[0][0], hi = matrix[n-1][n-1]; while (lo < hi) { int mid = lo + (hi - lo) / 2; int count = 0, r = n-1, c = 0; while (r >= 0 && c < n) { if (matrix[r][c] <= mid) { count += r+1; c++; } else r--; } if (count < k) lo = mid+1; else hi = mid; } return lo; }. Uses lo + (hi-lo)/2 instead of (lo+hi)/2 to avoid integer overflow when matrix values are near Integer.MAX_VALUE.

Performance comparison: binary search is O(n log(max-min)) time with O(1) space; heap is O(k log n) time with O(n) space. Binary search is better when k is large (say k = n*n/2) because the heap must perform k pops. The heap is better when k is tiny (k = 1 to 10) and n is large, because the binary search still runs O(n log(max-min)) iterations regardless of k. For most competitive programming cases and interview questions, the binary search is preferred because it uses constant space and has a cleaner worst-case bound.

The kth smallest in sorted matrix problem is part of the selection family: problems that ask for the kth element in some ordered set without fully sorting it. The closest relative is Kth Largest Element in an Array (LeetCode 215), which uses quickselect — a partition-based selection algorithm that is O(n) average case. For the matrix problem, quickselect does not directly apply because the elements are not in a contiguous array, but the mental model of selection-without-sorting is the same.

Find Median From Data Stream (LeetCode 295) is another selection problem: the median is the n/2-th element, and the heap-based solution using a max-heap and min-heap of equal sizes is structurally related to the min-heap row-merge approach for this problem. Search a 2D Matrix (LeetCode 74) and Search a 2D Matrix II (LeetCode 240) also use the staircase traversal technique — start from a corner where movement in each direction has a known monotone effect — making them natural companions to study alongside LeetCode 378.

The broader lesson is that sorted structure enables fast counting and searching without exhaustive enumeration. Whether it is a sorted array enabling binary search, a sorted matrix enabling staircase traversal, or a heap enabling O(log n) minimum extraction, the core principle is the same: exploit ordering to make definitive decisions at each step and eliminate large portions of the search space. This principle underlies nearly every efficient algorithm for ordered selection queries.