[LintCode] 401 Kth Smallest Number in Sorted Matrix 解题报告

Description
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 
Solve it in O(k log n) time where n is the bigger one between row size and column size.


思路
k路归并。
开一个min-heap，把每个array的第一个元素扔进去。然后每次拿出最小的，一直到拿第K次。
每拿出一个，要再放进去一个，就是拿出来这个所对应的array中的下一个元素要放进堆。
复杂度：
堆的add的复杂度是O(logn)，poll是O(1)。我们要搞k次，就是O(klogn)

Code

public class Solution {
    /**
     * @param matrix: a matrix of integers
     * @param k: an integer
     * @return: the kth smallest number in the matrix
     */
     
    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;
        }
    }
    public int kthSmallest(int[][] matrix, int k) {
        // write your code here
        
        Queue<Node> heap = new PriorityQueue<Node>(k, new Comparator<Node>() {
            public int compare(Node a, Node b) {
                return a.val - b.val;
            }
        });

        for (int i = 0; i < matrix.length; i++) {
            heap.add(new Node(i, 0, matrix[i][0]));
        }
        
        for (int i = 0; i < k; i++) {
            Node tmp = heap.poll();
            int row;
            int col;
            int val;
            row = tmp.row;
            col = tmp.col;
            val = tmp.val;
            
            if (i == k - 1) {
                return val;
            }
            
            if (col + 1 < matrix[row].length) {
                heap.add(new Node(row, col + 1, matrix[row][col + 1]));
            }
        }
        
        return -1;
    }
}