378. Kth Smallest Number in Sorted Matrix

Find the kth smallest number in at row and column sorted matrix.

Example

Given k = 4 and a matrix:

[
  [1 ,5 ,7],
  [3 ,7 ,8],
  [4 ,8 ,9],
]

return 5.

分析：

Although the matrix is horizontally and vertically sorted, there is no rule to find the next smallest number. In this solution, we will put the numbers in the first row to a heap, and remove the smallest one and add a new number whose position is right under the removed one. 

Note: when requesting to find the kth largest/smallest number, you should consider heap sort.

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        // be aware, the matrix may not be full matrix.
        if (matrix == null || matrix.length == 0 || matrix[0].length == 0 || k > matrix.length * matrix[0].length) {
            throw new IllegalArgumentException("Invalid input");
        }
        return helper(matrix, k);
    }
    
    private int helper(int[][] matrix, int k) {
        Queue<Point> heap = new PriorityQueue<>(k, (p1, p2) -> matrix[p1.row][p1.col] - matrix[p2.row][p2.col]);
        for (int i = 0; i < Math.min(matrix[0].length, k); i++) {
            heap.offer(new Point(0, i));
        }
        for (int i = 1; i <= k -1; i++) {
            Point curr = heap.poll();
            if (curr.row + 1 < matrix.length) {
                heap.offer(new Point(curr.row + 1, curr.col));
            }
        }
        return matrix[heap.peek().row][heap.peek().col];
    }
}

class Point {
    public int row, col;
    public Point(int row, int col) {
        this.row = row;
        this.col = col;
    }
}