Lintcode 543. Kth Largest in N Arrays

Find K-th largest element in N arrays.

Example

Example 1:

Input:
k=3, [[9,3,2,4,7],[1,2,3,4,8]]
Output:
7
Explanation:
the 3rd largest element is 7

Example 2:

Input:
k = 2, [[9,3,2,4,8],[1,2,3,4,2]]
Output:
8
Explanation:
the 1st largest element is 9, 2nd largest element is 8, 3rd largest element is 4 and etc.

Notice

You can swap elements in the array

Code (Java):

public class Solution {
    /**
     * @param arrays: a list of array
     * @param k: An integer
     * @return: an integer, K-th largest element in N arrays
     */
    public int KthInArrays(int[][] arrays, int k) {
        // sort it each array first//
        //
        for (int i = 0; i < arrays.length; i++) {
            Arrays.sort(arrays[i]);
        }

        // step 2: put the head of each array into a max heap
        //
        PriorityQueue<Pair> pq = new PriorityQueue<>(new MyPQComparator());

        for (int i = 0; i < arrays.length; i++) {
            if (arrays[i] != null && arrays[i].length > 0) {
                pq.offer(new Pair(i, arrays[i].length - 1, arrays[i][arrays[i].length - 1]));
            }
        }

        // step 3: poll k elements from the pq
        //
        int ans = -1;
        for (int i = 0; i < k; i++) {
            Pair curr = pq.poll();
            if (i == k - 1) {
                ans = curr.val;
            }

            if (curr.col - 1 >= 0) {
                pq.offer(new Pair(curr.row, curr.col - 1, arrays[curr.row][curr.col - 1]));
            }
        }

        return ans;
    }
}

class Pair {
    public int row;
    public int col;
    public int val;
    public Pair(int row, int col, int val) {
        this.row = row;
        this.col = col;
        this.val = val;
    }
}

class MyPQComparator implements Comparator<Pair> {
    @Override
    public int compare(Pair a, Pair b) {
        return b.val - a.val;
    }
}