Kth Smallest Number in Sorted Matrix - Lintcode

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

Example

Given k =4and a matrix:

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

Analysis

As soon as the question is asking for a smallest number, heap should be considered. And when it mentions about "sorted", then the complete traversal should not be required. The key point for this question is that for each point in a sorted matrix, the next larger number in the matrix is always its right or below neighbor.

Solution

class Pair {
    int x, y, val;
    public Pair(int x, int y, int val){
        this.x = x;
        this.y = y;
        this.val = val;
    }
}


public int kthSmallest(int[][] matrix, int k) {
    if (matrix.length == 0 || matrix[0].length == 0 || k < 0){
        return -1;    
    }

    PriorityQueue<Pair> heap = new PriorityQueue<Pair>(k, new Comparator<Pair>(){
        public int compare(Pair p1, Pair p2){
            return p1.val - p2.val;
        }
    });

    int[] deltaX = new int[]{0, 1};
    int[] deltaY = new int[]{1, 0};
    heap.offer(new Pair(0,0,matrix[0][0]));
    int[][] visited = new int[matrix.length][matrix[0].length];

    for (int i = 0; i < k - 1; i++){
        Pair p = heap.poll();
        for (int j = 0; j < 2; j++){
            int newX = p.x + deltaX[j];
            int newY = p.y + deltaY[j];
            if (newX < matrix.length && newY < matrix[0].length && visited[newX][newY] == 0){
                heap.offer(new Pair(newX, newY, matrix[newX][newY]));
                visited[newX][newY] = 1;
            }
        }      
    }

    return heap.poll().val;
}