Returning m greatest elements from k sorted array

Question

We have $k$ sorted arrays, $A_1[1...n_1],...,A_k[1..n_k]$, where $n_1+n_2+...+n_k=n$.

How can we get the $m$ greatest elements in running time $O(k + m\lg k)$?

I have tried to use MIN-HEAP size of $k$ since we have $k$ arrays. I have to navigate through the ends of the $k$ array somehow. By extracting the min of heap and getting new element from respective array was my idea but it does not work. I could not figure out getting m greatest elements.

Answer 1

Your idea works out. The idea is to maintain a heap of size $k$ which contains exactly one element from each array. For each array $A_i$, we maintain an index $p_i$ which keeps track of the position of the current element in the array. Initially, $p_i = n_i$, and we insert $A_i[n_i]$ into the heap. We then repeat the following $m$ times:

1. Extract the maximum element from the heap.
2. If the extracted element belonged to $A_i$, then we decrease $p_i$ by one and insert $A_i[p_i]$ into the heap.

The initialization phase takes time $O(k)$, and the extraction phase takes time $O(m\log k)$, for a total of $O(k+m\log k)$. You can prove that the extracted elements are indeed the $m$ largest ones.