Merging K- Sorted Lists using Priority Queue

Question

I have been asked in my algorithm class to make a K-way merge algorithm which is of O(nlogk) After searching i found it could be done via making a k length prio

Answer 1

Paul Hankin's solution is right, but it's a little bit difficult to read, especially you wanna implement in c++ or java. My solution is similar to Paul's. If you write in c++ or java, you may need an extra data structure to store the value of the element, index of the element in k-th array and index of the array in lists.

Element{
    int value;
    int idInArray,
    int idInList
}

But in python, I just store in a tuple (value, idInArray, idInList)

def mergeKArray(*lists):
    # implemented by min heap
    h = []
    r = []
    for k, arr in enumerate(lists):
        heapq.heappush(h, (arr[0], 0, k))
    while h:
        # min is the minimum element
        # i is the index of the min in the k-th array
        # k is the index of array in the list
        min, i, k = heapq.heappop(h)
        r.append(min)
        if i < len(lists[k]) - 1:
            i += 1
            heapq.heappush(h, (lists[k][i], i, k))
    return r

Because I just need to maintain a min heap with k elements, the time of eject or inject heap is O(log k). I also need to scan all n elements, and each element cost 2*log(k) time to inject and eject. Therefore, the big-O is O(n*log k).