Kth Smallest Element in a Sorted Matrix
Medium · Heap / Priority Queue · Explanation and code example
Heap Merge of Sorted Matrix Rows
How it works
Treat each matrix row as a sorted stream. A min-heap containing the next unconsumed cell of each relevant row always exposes the globally next value. Remove k - 1 values, pushing the successor from each removed value’s row, then read the heap minimum.
Steps
- Seed the first cell of up to min(k, rows) rows.
- Pop and advance the same row k - 1 times.
- Return the current smallest heap value.
Complexity
- Time: O(k log(r + 1)), with r = min(k, rows).
- Space: O(r).
Things to watch for
- Duplicate values count as separate elements.
- k is one-based.
- Use the value to order the heap, rather than row or column indices.
Python code example
import heapq
class Solution:
def kthSmallest(self, matrix: list[list[int]], k: int) -> int:
minHeap = [] # (matrix[i][j], i, j)
i = 0
while i < k and i < len(matrix):
heapq.heappush(minHeap, (matrix[i][0], i, 0))
i += 1
while k > 1:
k -= 1
_, i, j = heapq.heappop(minHeap)
if j + 1 < len(matrix[0]):
heapq.heappush(minHeap, (matrix[i][j + 1], i, j + 1))
return minHeap[0][0]
C++ code example
struct T {
int i;
int j;
int num; // matrix[i][j]
};
class Solution {
public:
int kthSmallest(vector<vector<int>>& matrix, int k) {
auto compare = [&](const T& a, const T& b) { return a.num > b.num; };
priority_queue<T, vector<T>, decltype(compare)> minHeap(compare);
for (int i = 0; i < k && i < matrix.size(); ++i)
minHeap.emplace(i, 0, matrix[i][0]);
while (k-- > 1) {
const auto [i, j, _] = minHeap.top();
minHeap.pop();
if (j + 1 < matrix[0].size())
minHeap.emplace(i, j + 1, matrix[i][j + 1]);
}
return minHeap.top().num;
}
};
Source and license
Explanation by LC Tracker. Code adapted from walkccc / Peng-Yu Chen · MIT license. Node types are supplied by LeetCode where needed.