[LintCode] Kth Smallest Number in Sorted Matrix

Problem

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

Challenge

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

Note

Solution

I. Muggle (95% ac, last case exceeded time limit)

public class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int size = matrix.length * matrix[0].length;
        if (matrix == null) return 0;
        PriorityQueue<Integer> queue = new PriorityQueue<Integer>(size+1);
        for (int i = 0; i < matrix.length; i++) {
            for (int j = 0; j < matrix[0].length; j++) {
                queue.add(matrix[i][j]);
            }
        }
        for (int i = 1; i < k; i++) {
            queue.remove();
        }
        return queue.peek();
    }
}

II. 堆排序

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]];
    }
}