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

1-kth-smallest-in-sorted-matrix.js

export function kthSmallestNumber(matrix, k) {
  const heap = new MinHeap();

  // initialize by pushing first element of each row into heap
  // since rows and cols are sorted, these are first n smallest elements in the matrix,
  // where n is the number of rows / cols
  for (let i = 0; i < matrix.length; ++i) {
    heap.offer([matrix[i][0], i, 0]);
  }

  // track how many elements we have popped from the heap,
  // we need to pop exactly k elements, or less if k > n*n
  let count = 0;
  // the last popped element will be here
  let popped = null;

  // stop when we have popped k element or the heap is empty
  while (count < k && heap.size() > 0) {
    const [value, row, col] = heap.poll();
    popped = value;
    count += 1;
    // get the next element in the row our popped element is coming from
    // and put it into the heap
    if (matrix[row][col + 1] != null) {
      heap.offer([matrix[row][col + 1], row, col + 1]);
    }
  }

  return popped;
}

2-README.md