Kth Smallest Element in a Sorted Matrix

medium matrix heap

Problem

Given an n×n matrix where rows and columns are sorted ascending, return the k-th smallest element. Memory must be sub-linear in n².

Inputmatrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8

Output13

Push the first column into a min-heap; on each pop, push the cell to the right.

matrix (JSON, sorted rows & cols)

import heapq

def kth_smallest(matrix, k):
    n = len(matrix)
    heap = [(matrix[i][0], i, 0) for i in range(min(n, k))]
    heapq.heapify(heap)
    for _ in range(k - 1):
        v, r, c = heapq.heappop(heap)
        if c + 1 < n:
            heapq.heappush(heap, (matrix[r][c + 1], r, c + 1))
    return heap[0][0]

// Simple min-heap by value.
function kthSmallest(matrix, k) {
  const n = matrix.length;
  const heap = [];
  for (let i = 0; i < Math.min(n, k); i++) heap.push([matrix[i][0], i, 0]);
  heap.sort((a, b) => a[0] - b[0]);
  for (let s = 0; s < k - 1; s++) {
    const [, r, c] = heap.shift();
    if (c + 1 < n) {
      const item = [matrix[r][c + 1], r, c + 1];
      let lo = 0, hi = heap.length;
      while (lo < hi) { const m = (lo + hi) >> 1; if (heap[m][0] < item[0]) lo = m + 1; else hi = m; }
      heap.splice(lo, 0, item);
    }
  }
  return heap[0][0];
}

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        PriorityQueue<int[]> heap = new PriorityQueue<>((a, b) -> a[0] - b[0]);
        for (int i = 0; i < Math.min(n, k); i++) heap.offer(new int[]{matrix[i][0], i, 0});
        for (int s = 0; s < k - 1; s++) {
            int[] top = heap.poll();
            if (top[2] + 1 < n) heap.offer(new int[]{matrix[top[1]][top[2] + 1], top[1], top[2] + 1});
        }
        return heap.peek()[0];
    }
}

int kthSmallest(vector<vector<int>>& matrix, int k) {
    int n = matrix.size();
    priority_queue<tuple<int,int,int>, vector<tuple<int,int,int>>, greater<>> heap;
    for (int i = 0; i < min(n, k); i++) heap.push({matrix[i][0], i, 0});
    for (int s = 0; s < k - 1; s++) {
        auto [v, r, c] = heap.top(); heap.pop();
        if (c + 1 < n) heap.push({matrix[r][c + 1], r, c + 1});
    }
    return get<0>(heap.top());
}

Explanation

The matrix has every row sorted left to right, so the smallest value in each row is its first column. We use a min-heap to pull values out in increasing order, like merging sorted lists.

We seed the heap with the first column: one entry (value, row, col=0) per row. Each entry remembers where it came from so we know which value comes next in that row.

Then we pop the smallest k-1 times. Every time we pop a cell, we push its right neighbour matrix[r][c+1] (if it exists), because that is the next-smallest candidate from that row. This keeps a moving "wavefront" of the smallest unseen value per row.

After k-1 pops, the value now at the heap top is the k-th smallest, and we return it.

Example: matrix=[[1,5,9],[10,11,13],[12,13,15]], k=8. Popping in order yields 1, 5, 9, 10, 11, 12, 13, and the 8th pop reaches 13. The heap only ever holds about one entry per row, so memory stays small.

Time: O(k log min(n, k)) Space: O(min(n, k))