Kth Smallest Number in Sorted Matrix

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.

Solution:

Similar to Merge k Sorted Arrays. SInce the matrix is sorted, the smallest value is always at the top left. And the second smallest should be next to it. In the example, the smallest is 1 (0,0 position), the second smallest is either 3 (1, 0 position) or 5 (0, 1 position). Then the third smallest will be the number adjecent to the smallest and second smallest.

To achieve this, we need to use a priorityQueue. We always pop out the minimal number in the queue and treat as the i-th minimal, and then we add its neighbor back to our min queue. However, there may be some duplicats (or revisiting), for example, 7 in (1,1) is neighboring with both 3 (1,0) and 5 (0, 1). Thus, we have to use a HashSet to mark the index of all visited cells by i * n + j. O(k log n) with space: O(k + n)

Code:

class Point {
    int val;
    int x;
    int y;
    public Point(int inVal, int inX, int inY) {
        this.val = inVal;
        this.x = inX;
        this.y = inY;
    }
}

public class Solution {

    private Comparator<Point> pointComparator = new Comparator<Point>() {
        public int compare(Point a, Point b) {
            return a.val - b.val;
        }
    };

    public int kthSmallest(int[][] matrix, int k) {
        if (matrix == null || matrix.length == 0) {
            return Integer.MIN_VALUE;
        }
        int m = matrix.length;
        int n = matrix[0].length;
        Queue<Point> q = new PriorityQueue(Math.min(m, n), pointComparator);
        HashSet<Integer> visited = new HashSet<>(); 
        kthSmallestHelper(0, 0, m, n, matrix, q, visited);
        int result = 0;
        int counter = 0;
        while (counter < k) {
            if (q.isEmpty()) {
                return Integer.MIN_VALUE;
            }
            Point p = q.poll();
            result = p.val;
            kthSmallestHelper(p.x + 1, p.y, m, n, matrix, q, visited);
            kthSmallestHelper(p.x, p.y + 1, m, n, matrix, q, visited);
            counter++;
        }
        return result;
    }

    private void kthSmallestHelper(int x, int y, int m, int n, int[][] matrix, 
                    Queue<Point> q, HashSet<Integer> visited) {
        if (x < m && y < n && !visited.contains(x * n + y)) {
            q.offer(new Point(matrix[x][y], x, y));
            visited.add(x * n + y);
        }    
    }
}