Function to merge $k$ sorted lists. This algorithm uses a queue data structure. The keys into the queue are the first element (the head) of each sorted list, and the values are the remaining elements of the list (the tail). Using a (min) queue guarantees that the list with the smallest first element will always be used first. The first for loop runs in $O(k*log(k))$ time, since the heads of $k$ lists must be heapified; the while loop runs in $O(k*n*log(k))$ time, where there are $k$ lists, and assuming all $k$ lists have $n$ elements.
import heapq
def merge_k_sorted_lists(lists):
h = []
for l in lists:
head = l[0]
l.pop(0)
rest = l
h.append((head, rest))
heapq.heapify(h)
r = []
while h != []:
head, rest = heapq.heappop(h)
r.append(head)
if rest != []:
head = rest.pop(0)
heapq.heappush(h, (head, rest))
return r
lists = [range(5), range(-5,1), range(5,9)]
print lists
listsmerged = merge_k_sorted_lists(lists)
print listsmerged
No comments:
Post a Comment