lintcode排序矩阵中的从小到大第k个数-401(Java版)

在一个排序矩阵中找从小到大的第 k 个整数。

排序矩阵的定义为：每一行递增，每一列也递增。

样例

给出 k = 4 和一个排序矩阵：

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

返回 5。

挑战

使用O(k log n)的方法，n为矩阵的宽度和高度中的最大值。

public class Solution {
    public int kthSmallest(final int[][] matrix, int k) {
        boolean[][] visited = new boolean[matrix.length][matrix[0].length];
        PriorityQueue<int[]> heap = new PriorityQueue<int[]>(k, new Comparator<int[]>(){
            public int compare(int[] a, int[] b) {
                return Integer.compare(matrix[a[0]][a[1]], matrix[b[0]][b[1]]);
            }
        });
        heap.add(new int[]{0,0});
        visited[0][0] = true;
        while (k > 1) {
            int[] res = heap.remove();
            int x = res[0];
            int y = res[1];
             
             
            if (x+1 < matrix.length && visited[x+1][y] == false) {
                visited[x+1][y] = true;
                heap.add(new int[]{x+1, y});
            }
            if (y+1 < matrix[0].length && visited[x][y+1] == false) {
                visited[x][y+1] = true;
                heap.add(new int[]{x, y+1});
            }
            k--;
        }
        int[] res = heap.remove();
        return matrix[res[0]][res[1]];
    }
}