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

Problems
21

Same as previous with a small tweak:

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) {
        min_heap.push(col[0]);
    }
    int idx[matrix.size()] = {};
    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 (!--k) return min_heap.top();
                min_heap.pop();
                if (++idx[row] < matrix[row].size()) min_heap.push(matrix[row][idx[row]]);
            }
        }
    }

    return 0;
}
22

pushing everything into a min heap and popping would make the heap size n and popping would be O(log(N)).

the heap size is k in the solution approach and popping is O(log(k)), which is faster

23

creating a min heap using all the elements would make the heap size n, and popping would be O(log(N)).

the heap size is k in the solution approach and popping is O(log(k)), which is faster

24

I completely agree. If for some reason there is something wrong with approaching this problem like the merging K sorted lists problem I think that should be explained.

25

+1

Did it the same way as Merge K sorted lists, it’s much more intuitive. The code in the solution is hard to follow.

function kthSmallest(matrix, k) {
    const pq = new MinHeap()
    // add first element from each row
    for (let row = 0; row < matrix.length; row++) pq.push(
        new HeapItem([row, 0], matrix[row][0])
    )

    // pop + push k-1 times similarly to merge-k-lists
    while (--k) {
        let [row, col] = pq.pop().item
        col++
        let item = new HeapItem( [row, col], matrix[row][col] )
        if (col < matrix[row].length) pq.push(item)
    }

    // return the kth element
    const [row, col] = pq.pop().item
    return matrix[row][col]
}
26

better understanding version. using k-way merge.

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

# min heap
min_heap = [] # (value, row, column)

# push every row first column into heap
for row in range(n):
    heappush(min_heap, (matrix[row][0], row, 0))

res = -1

# key-way merge, and find the k-th element.
while min_heap and k > 0:
    val, row, col = heappop(min_heap)
    res = val # update result
    k -= 1
    
    # push the next node into heap
    if col + 1 < n:
        heappush(min_heap, (matrix[row][col + 1], row, col + 1))
    
return res
27

Easy problem once you realise that each sublist in matrix boils down to merging K lists. Python3 Solution :

    heap=[]
    x=0
    res=0
    n=len(matrix[0])
    for i in range(n):
        heappush(heap,(matrix[i][0],matrix[i],0))
    
    while x != k:
        
        res,subarray,ptr=heappop(heap)
        x+=1
        
        if(ptr+1==n):
            continue
        else:
            heappush(heap,(subarray[ptr+1],subarray,ptr+1))
    
    return res
28

[LeetCode - The World's Leading Online Programming Learning Platform]
Liked this explantion with diagram

29

My simple solution, using Alogmonster’s intuition:

from typing import List
import heapq
def kth_smallest(matrix: List[List[int]], k: int) -> int: 
    hp = []
    
    # adding initial pointers to the heap
    for row in range(len(matrix)):
        heapq.heappush(hp, (matrix[row][0],row,0))
    
    col_length = len(matrix[0])
    row_length = len(matrix)
    for counter in range(k):
        val, row, col = heapq.heappop(hp)
        
        if col + 1 < col_length and row < row_length: 
            heapq.heappush(hp, (matrix[row][col + 1], row, col + 1))
    return matrix[row][col]