Kth Smallest Element in a Sorted Matrix

Problem

Given a n x n matrix where each of the rows and columns are sorted in ascending order, find the kth smallest element in the matrix.

Note that it is the kth smallest element in the sorted order, not the kth distinct element.

Example:

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

return 13.

Note: You may assume k is always valid, 1 ≤ k ≤ n^2.

Solution

Priority Queue Solution: O(n^2) time, O(n) space

public class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        if (matrix == null || matrix.length == 0) return 0;
        Queue<Integer> q = new PriorityQueue<>(Collections.reverseOrder());
        for (int i = 0; i < matrix.length; i++) {
            for (int j = 0; j < matrix[0].length; j++) {
                int num = matrix[i][j];
                if (q.size() < k) q.offer(num);
                else if (num < q.peek()) {
                    q.poll();
                    q.offer(num);
                }
            }
        }
        return q.poll();
    }
}

Analysis

A simple solution using PriorityQueue
Because queue is FIFO, we init the q in reverse order
Then we loop through the matrix and as long as q.size() < k, we add it to q
After q.size() >= k, we only add current element if it's smaller than q.peek()
This is because we are trying to find nth smallest number At the end, we return q.poll()