Kth Smallest Element in a Sorted Matrix

Find the kth smallest element in an (n × n) matrix, where each row and column of the matrix is sorted in ascending order

Kth Smallest Element in a Sorted Matrix

Problem Statement

You are given an n × n matrix. Every row and every column in this matrix is sorted in ascending order.

Find the kth smallest element in the matrix.

Even if there are repeated values in the matrix, each element is still counted as its own unique entry when finding the kth smallest.

Constraints

• n == matrix.length
• n == matrix[i].length
• 1 ≤ n ≤ 100
• −10³ ≤ matrix[i][j] ≤ 10³
• 1 ≤ k ≤ n²

Examples

Example 1

Input:

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

Explanation: Flattening and sorting all elements gives [1, 5, 9, 10, 11, 12, 13, 13, 15]. The 8th smallest element is 13.

Output: 13

Example 2

Input:

matrix = [
  [-5]
]
k = 1

Explanation: The matrix has only one element, -5, which is trivially the 1st smallest element.

Output: -5

Example 3

Input:

matrix = [
  [1, 2],
  [1, 3]
]
k = 2

Explanation: Flattening and sorting all elements (keeping duplicates) gives [1, 1, 2, 3]. The 2nd smallest element is 1.

Output: 1

Solution

We use a min heap to solve this efficiently, without needing to flatten and sort the entire matrix.

Steps

• We insert the first element of every row into the heap. Along with each value, we also store its row index and column index.
• We repeatedly pop the smallest value from the heap:
  • Each time we pop a value, we count it as one more number checked.
  • If this is the kth value we've checked, we stop here — this is our answer.
  • Otherwise, we move to the next column in the same row this value came from (if one exists), and push that next value into the heap.
• Since every row and column is sorted, this approach always finds the next-smallest value without needing to look at the whole matrix at once.

Code

function kthSmallestNumber(matrix, k) {
    let rowCount = matrix.length,
        minNumbers = new MinHeap();

    for (let index = 0; index < rowCount; index++) {
        minNumbers.offer([matrix[index][0], index, 0]);
    }

    let numbersChecked = 0,
        smallestElement = 0,
        rowIndex, colIndex;

    while (minNumbers.size() > 0) {
        let result = minNumbers.poll();
        [smallestElement, rowIndex, colIndex] = result;
        numbersChecked += 1;

        if (numbersChecked == k) break;

        if (colIndex + 1 < matrix[rowIndex].length) {
            minNumbers.offer([matrix[rowIndex][colIndex + 1], rowIndex, colIndex + 1]);
        }
    }

    return smallestElement;
}