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.

public class Solution {

    /**

     * @param matrix: a matrix of integers

     * @param k: an integer

     * @return: the kth smallest number in the matrix

     */

    public int kthSmallest(final int[][] matrix, int k) {

        // write your code here

        int m = matrix.length, n = matrix[0].length;

        boolean[][] visited = new boolean[m][n];

        

        PriorityQueue<int[]> queue = new PriorityQueue<int[]>(k, new Comparator<int[]>(){

            public int compare(int[] a, int[]b){

                return matrix[a[0]][a[1]] - matrix[b[0]][b[1]];

            }

        });

        queue.add(new int[]{0,0});

        visited[0][0] = true;

        while(k > 1){

            int[] res = queue.poll();

            k--;

            int x = res[0];

            int y = res[1];

            if(x+1 < m && visited[x+1][y] == false){

                queue.add(new int[]{x+1, y});

                visited[x+1][y] = true;

            }

            

            if(y+1 < n && visited[x][y+1] == false){

                queue.add(new int[]{x, y+1});

                visited[x][y+1] = true;

            }

        }

        int[] res = queue.poll();

        return matrix[res[0]][res[1]];

    }

}