[LintCode] 543 Kth Largest in N Arrays 解题报告

Description
Find K-th largest element in N arrays.

Notice
You can swap elements in the array


Example
In n=2 arrays [[9,3,2,4,7],[1,2,3,4,8]], the 3rd largest element is 7.

In n=2 arrays [[9,3,2,4,8],[1,2,3,4,2]], the 1st largest element is 9, 2nd largest element is 8, 3rd largest element is 4 and etc.

思路
k路归并。
先把所有的array都排个序。
然后开一个堆，把每个array的最后一个数丢进去。
这里要注意，我们要自己写一个Comparator实现max-heap。
然后开始找第K个最大。每次从堆里面拿出来一个（一定是最大的）。
如果是第K个，那么就找到了。如果不是，从这个拿出来的数所在的array，把倒数第二个数字丢进堆里。

Code

public class Solution {
    /**
     * @param arrays a list of array
     * @param k an integer
     * @return an integer, K-th largest element in N arrays
     */
     
    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 KthInArrays(int[][] arrays, int k) {
        // Write your code here

        Queue<Node> heap = new PriorityQueue<Node>(k, new Comparator<Node>() {
           public int compare(Node a, Node b) {
               if (a.val > b.val) {
                   return -1;
               }
               else if (a.val < b.val) {
                   return 1;
               }
               return 0;
           } 
        });
        
        for (int i = 0; i < arrays.length; i++) {
            Arrays.sort(arrays[i]);
            
            if (arrays[i].length > 0) {
                int row = i;
                int col = arrays[i].length - 1;
                int val = arrays[row][col];
                heap.add(new Node(row, col, val));
            }
        }
        
        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 tmp.val;
            }
            
            if (col - 1 >= 0) {
                heap.add(new Node(row, col - 1, arrays[row][col - 1]));
            }
            
        }    
        
        return -1;
    }
}