Lintcode: 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

KLog(Min(M,N,K))时间复杂度 K是因为要Poll K次并且同时insert K次，Min(M,N,K)是堆的size，insert的时间是Log(MIN(M,N,K))

 

思路就是维护一个最小堆，先把第一行，或第一列（本题做法是第一列，之后默认第一列）加入堆中，poll K次，每poll一次之后要把poll的元素的下一个元素加入堆中，本题就是poll的元素的下一列元素。最后一次poll的元素即为所求。因为需要知道每个poll的元素的位置，所以写了个Point Class

public class Solution {
    /**
     * @param matrix: a matrix of integers
     * @param k: an integer
     * @return: the kth smallest number in the matrix
     */
    
        // write your code here
        public class Point {
            public int x, y, val;
            public Point(int x, int y, int val) {
                this.x = x;
                this.y = y;
                this.val = val;
            }
        } 
        
        Comparator<Point> comp = new Comparator<Point>() {
            public int compare(Point left, Point right) {
                return left.val - right.val;
            }
        };
        
        public int kthSmallest(int[][] matrix, int k) {
            if (matrix == null || matrix.length == 0 || matrix[0].length == 0) {
                return 0;
            }
            if (k > matrix.length * matrix[0].length) {
                return 0;
            }
            return horizontal(matrix, k);
        }
        
        private int horizontal(int[][] matrix, int k) {
            Queue<Point> heap = new PriorityQueue<Point>(k, comp);
            for (int i = 0; i < Math.min(matrix.length, k); i++) {
                heap.offer(new Point(i, 0, matrix[i][0]));
            }
            for (int i = 0; i < k - 1; i++) {
                Point curr = heap.poll();
                if (curr.y + 1 < matrix[0].length) {
                    heap.offer(new Point(curr.x, curr.y + 1, matrix[curr.x][curr.y + 1]));
                }
            }
            return heap.peek().val;
        }
            
        
}