Kth Smallest Element in a Sorted Matrix
Understanding the Problem
We have a square matrix where each row is sorted left to right and each column is sorted top to bottom. We need to find the kth smallest element overall. The complication is that while rows and columns are individually sorted, elements across different rows can interleave. The last element of row 0 might be larger than the first element of row 2, so there is no single fixed ordering we can read off the matrix.
The matrix does have useful structure to lean on. matrix[0][0] is always the global minimum and matrix[n-1][n-1] is always the global maximum, and any path that moves only right or down passes through non-decreasing values. That makes the matrix behave like n sorted lists stacked together, which suggests two angles: merge those lists with a heap, or binary search on the answer value itself.
Key Constraints:
1 <= n <= 300→ The matrix holds up to 90,000 elements, so time is not a concern, but the problem requires better than O(n^2) memory. That rules out any solution that copies the whole matrix into another structure.-10^9 <= matrix[i][j] <= 10^9→ Values fit in a 32-bit signed integer, andmatrix[0][0]tomatrix[n-1][n-1]bounds the answer, so binary search on the value runs in about 31 iterations regardless of n.All rows and columns sorted→ Each row is a sorted list we can merge with a heap, and the row-and-column ordering lets us count elements<= vin O(n) time with a staircase walk.
Approach 1: Sort All Elements
Intuition
Copy every element into a flat list, sort it, and return the element at index k - 1. This ignores the sorted structure of the matrix, but it is correct by construction and gives us a baseline to improve on.
Algorithm
- Create an empty list.
- Iterate through every element in the matrix and add it to the list.
- Sort the list.
- Return the element at index
k - 1.
Example Walkthrough
Code
- Time Complexity: O(n^2 log n). We flatten all n^2 elements and sort them. Sorting dominates at O(n^2 log(n^2)) = O(n^2 log n).
- Space Complexity: O(n^2). We store all n^2 elements in a separate list.
This violates the O(n^2) memory limit, since the flat list holds a full copy of the matrix. The next approach uses the fact that each row is already sorted: it merges the rows with a min-heap and keeps only n elements in memory at a time.
Approach 2: Min-Heap (Merge K Sorted Rows)
Intuition
Treat the matrix as n sorted lists, one per row, and do a k-way merge with a min-heap. Push the first element of every row into the heap, then extract the minimum k times. Each time you extract an element from row i, push the next element in that row (if one exists). The heap holds at most n elements, so each push and pop is O(log n), and we stop after k extractions instead of merging all rows completely.
The heap holds the smallest unprocessed element from each row that still has elements left. Because each row is sorted, no element behind those candidates can be smaller than its row's candidate. So the heap's minimum is the smallest element not yet extracted across the whole matrix. Popping it and pushing the next element from its row preserves that property, and the value popped on the kth iteration is the kth smallest overall.
Algorithm
- Create a min-heap. Push the first element of each row as a tuple:
(value, row, col). - Repeat k times:
- Pop the smallest element from the heap.
- If the popped element has a next element in its row (col + 1 < n), push that next element.
- The kth popped element is the answer.
Example Walkthrough
Code
- Time Complexity: O(k log n). We initialize the heap with n elements, then perform k extract-and-insert operations, each costing O(log n).
- Space Complexity: O(n). The heap holds at most n elements at any time (one per row).
The heap cost grows with k. When k is close to n^2, this does nearly n^2 heap operations. The next approach removes the dependence on k by binary searching on the answer value and using both the row and column ordering to count quickly.
Approach 3: Binary Search on Value
Intuition
Instead of extracting elements in order, binary search on the answer value. The search range is [matrix[0][0], matrix[n-1][n-1]], from the smallest to the largest value in the matrix.
For a candidate value mid, count how many elements are <= mid. If that count is at least k, the kth smallest is mid or something smaller, so search the lower half by setting high = mid. If the count is less than k, the answer must be larger, so set low = mid + 1. The loop narrows low and high until they meet, and the converged value is the answer.
The counting step is where the row-and-column ordering pays off. Start at the bottom-left corner. If the current element is <= mid, then every element above it in the same column is also <= mid (the column is sorted), so add row + 1 to the count and move right to the next column. If the current element is > mid, no element below it in this column can qualify either, so move up one row. Each step either advances the column or drops a row, so the walk visits at most 2n cells and counts all qualifying elements in O(n) time.
The converged value is always an element of the matrix. The binary search returns the smallest value v for which count(<= v) >= k. If v were not in the matrix, then count(<= v) would equal count(<= v-1), so v - 1 would also satisfy the condition, contradicting that v is the smallest such value. So v must be a real matrix element, and it is exactly the kth smallest.
Algorithm
- Set
low = matrix[0][0]andhigh = matrix[n-1][n-1]. - While
low < high: - Compute
mid = low + (high - low) / 2. - Count how many elements are <= mid using the staircase traversal.
- If count >= k, set
high = mid(answer could be mid or smaller). - If count < k, set
low = mid + 1(answer must be larger). - Return
low.
Example Walkthrough
Code
- Time Complexity: O(n * log(max - min)). The binary search runs O(log(max - min)) iterations. Each iteration does a staircase count in O(n). With values up to 10^9, log(max - min) is at most ~63.
- Space Complexity: O(1). Only a handful of integer variables. No extra data structures.