Kth Smallest Element in a Sorted Matrix - Priority Queue / Heap / Top K


This question feels really similar to Merge K Sorted Lists. In fact, I used the same method to solve it.

Agreed. Problem reduces to Merge K Sorted Lists since each row is a sorted list, and we have N (# of cols) rows.

Current solution grows from matrix[0][0], which requires row-wise growth as well. I don’t see any benefit to this approach vs. seeding the heap with row[0] for all rows (solution to Merge K Sorted Lists).


anyone who is having trouble understanding the solution given above can refer to link mentioned, approach is similar but implementation is cleaner.

A detailed explanation of the solution would help.
A good klogk solution given here:

I found the leetcode explanation much clearer and implemented a js solution based on their python solution.

    const n = matrix.length;
    const minHeap = new MinHeap();
    let res;
    for (let row = 0; row < Math.min(n, k); row++) {
        const node = new HeapItem([matrix[row][0], row, 0], matrix[row][0]);
    while (k > 0) {
        const node = minHeap.pop();
        const [element, row, col] = node.item;
        if (col < n - 1) {
            const nextVal = matrix[row][col + 1];
            const nextItem = new HeapItem([nextVal, row, col + 1], nextVal);
        res = element;
    return res;

why don’t use the merge K sorted lists algorithm?

def kth_smallest(matrix: List[List[int]], k: int) → int:

min_heap = []
for rows in matrix:
    heappush(min_heap, (rows[0], rows, 0))

count = 0
res = None
while count < k:
    # we pop element from a row
    val, row, index = heappop(min_heap)
    res = val
    index += 1
    if index < len(row): 
        # push the next index's value from the row poped above
        heappush(min_heap, (row[index], row, index))
    count += 1

return res

Thanks for the link, this is easier to understand and implement.

I think a lot of the the solutions are not taking advantage of heapify and using heappush.

from heapq import heapify, heappop

def kth_smallest(matrix: List[List[int]], k: int) → int:
heap = []
for r in matrix:

for i in range(1, k):

return heappop(heap)

I believe there is a mistake on the first deck of slides. At K=6 it says “13 is the smallest one”, but the 3rd pointer points at 12 at this time.
The second deck doesn’t have that mistake.

was wondering the same. merging K sorted array approach is way more intuitive than offered above.

Below is my solution. I have made the variable names very clear, which follows Google’s coding conventions:

def kth_smallest(matrix: List[List[int]], k: int) → int:
heap = []
num_items_so_far = 0

for i in range(len(matrix[0])):
    # Get first element, matrice index, and element index for each sorted row
    heappush(heap, (matrix[i][0], i, 0))

while heap:
    val, matrix_index, item_index = heappop(heap)
    item_index += 1
    num_items_so_far += 1
    if num_items_so_far == k:
        return val
    current_matrix = matrix[matrix_index]
    if item_index < len(matrix[0]):
        heappush(heap, (current_matrix[item_index], matrix_index, item_index))

return None

Simple Soln

using namespace std;

using Pair = pair<int, int>;

string pairString(int x, int y) {
    return to_string(x) + to_string(y);

int kth_smallest(std::vector<std::vector<int>> matrix, int k) {
    int n = matrix.size();
    auto compareLambda = [&matrix](Pair pos1, Pair pos2) {
        return matrix[pos1.first][pos1.second] > matrix[pos2.first][pos2.second];
    priority_queue<Pair, vector<Pair>, decltype(compareLambda)> pq(compareLambda);
    pq.push(Pair(0, 0));
    unordered_set<string> visited;
    while(k-- > 1){
        Pair node = pq.top();
        int row = node.first, col = node.second;
        // Add the item on the right to the heap if everything above it is processed
        if (col + 1 < n && visited.count(pairString(row, col + 1)) == 0) {
            pq.push(Pair(row, col + 1));
            visited.insert(pairString(row, col + 1));
        // Add the item below it to the heap if everything before it is processed
        if (row + 1 < n && visited.count(pairString(row+1, col)) == 0) {
            pq.push(Pair(row + 1, col));
            visited.insert(pairString(row+1, col));
    Pair res = pq.top();
    return matrix[res.first][res.second];

I don’t understand the difference in the solution here to Kth Largest Element in an Array.

Why can’t we apply the same solution here with a minheap instead?

The solution here is kinda overcomplicated, I still can’t get why we need additional space column_top, row_first.
This is simplified, but honestly, I came up with my solution when I cheated and scrolled to the solution (after I solved it with transform matrix to flat list and then heapify and etc…)

def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
    ROWS, COLS = len(matrix), len(matrix[0])
    minHeap = [] # val : row : col
    for i in range(len(matrix)):
        heapq.heappush(minHeap, (matrix[i][0], i, 0))
    result = 0
    while k > 0:
        val, row, col = heapq.heappop(minHeap)
        result = val
        if col + 1 < COLS:
            heapq.heappush(minHeap, (matrix[row][col + 1], row, col + 1))

        k -= 1
    return result

I have the same question…

I think my solution runs in O(K)…
In my solution, the heap size is at most 3. In each iteration of the loop, pop the min element and push right cell and down cell to the heap, repeat K times. So it runs in K*log3. log3 is constant so runtime is O(K).

from typing import List
from heapq import heappush, heappop

def kth_smallest(matrix: List[List[int]], k: int) → int:
m = len(matrix)
n = len(matrix[0])
h = []
heappush(h, (matrix[0][0], (0, 0)))
matrix[0][0] = -1

i = 0 
while(i < k-1):
    _, (r, c) = heappop(h)
    i +=1 
    for dx, dy in [(0, 1), (1, 0)]:
        newx = r + dx
        newy = c + dy
        if newx < m and newy < n and matrix[newx][newy] != -1:
            heappush(h, (matrix[newx][newy], (newx, newy)))
            matrix[newx][newy] = -1
return heappop(h)[0]

nvm, miss the point that heap size will grow if i pop once but push twice in the loop

def kth_smallest(matrix: List[List[int]], k: int) → int:
n = len(matrix)
heap = [(l[0], l[1:]) for l in matrix]

for i in range(k):
    h = heapq.heappop(heap)
    if len(h[1]) > 0:
        heapq.heappush(heap, (h[1][0], h[1][1:]))
return h[0]

My solution applies the same logic but using an array in indexes:

int kth_smallest(std::vector<std::vector<int>> matrix, int k) {
    std::priority_queue<int, std::vector<int>, std::greater<int>> min_heap;
    for (const auto &col : matrix) {
    int idx[matrix.size()] = {};
    int order = 1;
    while (min_heap.size()) {
        for (int row = 0; row < matrix.size(); ++row) {
            while (idx[row] < matrix[row].size() && matrix[row][idx[row]] == min_heap.top()) {
                if (order == k) return min_heap.top();
                if (++idx[row] < matrix[row].size()) min_heap.push(matrix[row][idx[row]]);

    return 0;