lintcode-medium-Kth Smallest Number in Sorted Matrix

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

 

Given k = 4 and a matrix:

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

return 5

 

public class Solution {
    /**
     * @param matrix: a matrix of integers
     * @param k: an integer
     * @return: the kth smallest number in the matrix
     */
    public int kthSmallest(int[][] matrix, int k) {
        // write your code here
        
        if(matrix == null || matrix.length == 0 || matrix[0] == null || matrix[0].length == 0)
            return 0;
        
        int m = matrix.length;
        int n = matrix[0].length;
        k--;
        
        if(k < 0 || k > m * n - 1)
            return 0;
        
        PriorityQueue<point> queue = new PriorityQueue<point>(n, new Comparator<point>(){
            public int compare(point p1, point p2){
                return p1.val - p2.val;
            }
            
        });
        
        for(int i = 0; i < n; i++){
            queue.offer(new point(0, i, matrix[0][i]));
        }
        
        for(int i = 0; i < k; i++){
            point temp = queue.poll();
            
            if(temp.row + 1 < m){
                queue.offer(new point(temp.row + 1, temp.col, matrix[temp.row + 1][temp.col]));
            }
        }
        
        return queue.poll().val;
    }

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