[Lintcode] Kth Smallest Number in Sorted Matrix

Kth Smallest Number in Sorted Matrix

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

Have you met this question in a real interview? 

Yes

Example

Given k = 4 and a matrix:

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

return 5

Challenge

O(k log n), n is the maximal number in width and height.

实际上是和丑数问题类似的一个问题，也是使用一个优先队列

import java.util.*;

class Node{
    public int row;
    public int col;
    public int value;
    public Node(int row,int col, int value){
        this.row = row;
        this.col = col;
        this.value = value;
    }
}
public class Solution {
    /**
     * @param matrix: a matrix of integers
     * @param k: an integer
     * @return: the kth smallest number in the matrix
     */
    private void set(int row,int col,boolean [][] visited, int [][]matrix, PriorityQueue<Node> heap){
        int m = visited.length;
        int n = visited[0].length;
        if(row<0||row>=m||col<0||col>=n||visited[row][col]) return;
        heap.offer(new Node(row,col,matrix[row][col]));
        visited[row][col]=true;
    }
    public int kthSmallest(int[][] matrix, int k) {
        // write your code here
        Comparator<Node> mycompare = new Comparator<Node>(){
            public int compare(Node n1,Node n2){
                return n1.value-n2.value;
            }
        };
        boolean [][] visited = new boolean[matrix.length][matrix[0].length];
        PriorityQueue<Node> heap = new PriorityQueue<Node>(k,mycompare);
        heap.offer(new Node(0,0,matrix[0][0]));
        int selected = 0;
        visited[0][0]=true;
        for(int i=0;i<k;i++){
            Node node = heap.poll();
            selected = node.value;
            int row = node.row;
            int col = node.col;
            
            set(row+1,col,visited,matrix,heap);//down
            set(row,col+1,visited,matrix,heap);//right
        }
        return selected;
    }
}