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

Here's some Python 2 code that does the merging.

import heapq

def addtoheap(h, i, it):
    try:
        heapq.heappush(h, (next(it), i))
    except StopIteration:
        pass

def mergek(*lists):
    its = map(iter, lists)
    h = []
    for i, it in enumerate(its):
        addtoheap(h, i, it)
    while h:
        v, i = heapq.heappop(h)
        addtoheap(h, i, its[i])
        yield v

for x in mergek([1, 3, 5], [2, 4, 6], [7, 8, 9], [10]):
    print x

Why is it O(n log k)? Well for each value removed, there's a heap pop and possibly a heap push (both of which are O(log k)). Since we remove n elements, it's O(n log k).