mathjax

Saturday, April 26, 2014

merging k sorted lists

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