Kth Largest in N Arrays 543

Question

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 7 and etc.

Solution

1. 将每一行从小到大排序

2. 将每一行最大（最后）的元素加入一个最大堆

3. 取出堆顶元素，若此时取出元素数量不足k个，则将取出元素的左边元素（如果存在）加入堆

4. 重复3直到取出k个元素为止

代码如下：

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 value;
        int row;
        int column;
        public Node(int row, int column, int value){
            this.value = value;
            this.row = row;
            this.column = column;
        }
    }

    public int KthInArrays(int[][] arrays, int k) {
        // Write your code here
        if(arrays == null || arrays.length == 0 || k <= 0){
            return -1;
        }

        for(int i = 0; i < arrays.length; i++){
            Arrays.sort(arrays[i]);
        }

        Queue<Node> queue = new PriorityQueue<Node>(arrays.length, new Comparator<Node>(){
            public int compare(Node a, Node b){
                return b.value - a.value;
            }
        });

        for(int i = 0; i < arrays.length; i++){
            if(arrays[i].length > 0){
                queue.offer(new Node(i, arrays[i].length - 1, arrays[i][arrays[i].length - 1]));
            }
        }

        for(int i = 0; i < k; i++){
            Node temp = queue.poll();
            if(i == k - 1){
                return temp.value;
            }
            if(temp.column > 0){
                queue.offer(new Node(temp.row, temp.column - 1, arrays[temp.row][temp.column - 1]));
            }
        }

        return -1;
    }
}