Saturday, July 23, 2016

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

Use a PriorityQueue. First put the first row / column into the queue. Then poll out the first value, and add the next value (next row / column after this) into the queue for k - 1 times.


 public int kthSmallest(int[][] matrix, int k) {
        PriorityQueue<Point> queue = new PriorityQueue<>();
        int rows = matrix.length;
        int cols = matrix[0].length;
        for (int i = 0; i < cols; i++) {
            queue.add(new Point(matrix[0][i], 0, i));
        }
        
        while (k > 1) {
            Point curr = queue.poll();
            if (curr.x < rows - 1) {
                int x = curr.x;
                int y = curr.y;
                queue.add(new Point(matrix[x + 1][y], x + 1, y));
            }
            k--;
        }
        return queue.poll().val;
    }
    
    private class Point implements Comparable<Point> {
        int x;
        int y;
        int val;
        
        public Point (int val, int x, int y) {
            this.val = val;
            this.x = x;
            this.y = y;
        }
        
        @Override
        public int compareTo (Point another) {
            return this.val - another.val;
        }
    }

No comments:

Post a Comment