Lintcode 401. Kth Smallest Number in Sorted Matrix

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

Each row and each column of the matrix is incremental.

Example

Example 1:

Input:
[
  [1 ,5 ,7],
  [3 ,7 ,8],
  [4 ,8 ,9],
]
k = 4
Output: 5

Example 2:

Input:
[
  [1, 2],
  [3, 4]
]
k = 3
Output: 3

Challenge

O*(klogn*) time, n is the maximum of the width and height of the matrix.

Solution:
The top left is the smallest number of the matrix, and bottom right is the biggest number in the matrix. Starting from the top left, use a pq to store the value and its index. For each number polled from the pq, insert its two neighbors into the pq since they are the candidates for the next smallest elements. 

Note that we need to use a visited[][] to note the visited node otherwise we may end with a loop.

Code (Java):

public class Solution {
    /**
     * @param matrix: a matrix of integers
     * @param k: An integer
     * @return: the kth smallest number in the matrix
     */
    public int kthSmallest(int[][] matrix, int k) {
        int rows = matrix.length;
        int cols = matrix[0].length;
        boolean[][] visited = new boolean[rows][cols];
        
        PriorityQueue<Node> pq = new PriorityQueue<>(new MyPQComparator());
        
        pq.offer(new Node(0, 0, matrix[0][0]));
        visited[0][0] = true;
        
        int[][] dir = new int[][]{{0, 1}, {1, 0}};
        
        int ans = 0;
        for (int i = 0; i < k; i++) {
            Node curr = pq.poll();
            if (i == k - 1) {
                ans = curr.val;
            }
            
            for (int j = 0; j < 2; j++) {
                int nextRow = curr.row + dir[j][0];
                int nextCol = curr.col + dir[j][1];
                
                if (nextRow < rows && nextCol < cols && !visited[nextRow][nextCol]) {
                    pq.offer(new Node(nextRow, nextCol, matrix[nextRow][nextCol]));
                    visited[nextRow][nextCol] = true;
                }
            }
        }
        
        return ans;
    }
}

class Node {
    int row;
    int col;
    int val;
    
    public Node (int row, int col, int val) {
        this.row = row;
        this.col = col;
        this.val = val;
    }
}

class MyPQComparator implements Comparator<Node> {
    @Override
    public int compare(Node a, Node b) {
        return a.val - b.val;
    }
}

Analysis:
Time complexity is O(klogk). That's because there are k operations, for each operation, we pop out one element, but insert two. So the total number of elements in the pq is k. So the overall time complexity is O(klogk).