Kth Smallest Element in a Sorted Matrix

Leetcode Problem 378: Kth Smallest Element in a Sorted Matrix

Problem Summary

Given an n x n matrix where each row and each column is already sorted in non-decreasing order, return the value that appears in position k when all matrix elements are viewed as one sorted sequence.

Repeated numbers still count as separate positions, so you are looking for the kth element in sorted order, not the kth distinct value.

matrix:
1   5   9
10 11  13
12 13  15

sorted view:
1, 5, 9, 10, 11, 12, 13, 13, 15

The problem also requires an approach that uses less than O(n^2) extra memory, so materializing and sorting the whole matrix is intentionally not the target strategy.

Constraints:

• The matrix is always square, so the number of rows and columns is the same and equal to n.

• n can range from 1 to 300.

• Matrix values can range from -10^9 to 10^9.

• Rows are sorted from left to right and columns are sorted from top to bottom, both in non-decreasing order.

• The requested rank k is always valid and falls between 1 and n^2 inclusive.

• -109 <= matrix[i][j] <= 109

• All the rows and columns of matrix are guaranteed to be sorted in non-decreasing order.

• 1 <= k <= n2

Follow up:

• Could you solve the problem with a constant memory (i.e., O(1) memory complexity)?

• Could you solve the problem in O(n) time complexity? The solution may be too advanced for an interview but you may find reading this paper fun.

Techniques

• Array
• Binary Search
• Sorting
• Heap (Priority Queue)
• Matrix

Solution

import { Heap } from "../heap"

interface Position {
    i: number
    j: number
}

function kthSmallest(matrix: number[][], k: number): number {
    const heap = new Heap<{ position: Position, value: number }>((a, b) => a.value - b.value)
    const visited = new Set<string>()

    heap.insert({ position: { i: 0, j: 0}, value: matrix[0][0] })
    visited.add("0,0")

    while (--k > 0 && heap.size() > 0) {
        const { position } = heap.extract()!

        if (position.j + 1 < matrix[0].length && !visited.has(`${position.i},${position.j + 1}`)) {
            heap.insert({ 
                value: matrix[position.i][position.j + 1], 
                position: { i: position.i, j: position.j + 1 } 
            })
            visited.add(`${position.i},${position.j + 1}`)
        }

        if (position.j === 0 && position.i + 1 < matrix.length && !visited.has(`${position.i + 1},0`)) {
            heap.insert({ 
                value: matrix[position.i + 1][0], 
                position: { i: position.i + 1, j: 0 } 
            })
            visited.add(`${position.i + 1},0`)
        }
    }

    return heap.extract()!.value
}