Merge k sorted lists in O(nlogk) time

Here’s a neat little algorithm written in python to merge k sorted lists in O(nlogk) time using a k-element min-heap:

# Lists is an array of sorted lists (arrays):
#   [ [...], [...], … ]
def ListMerge(Lists):
  # The number of elements awaiting merge in each list.
  sizes = [len(L) for L in Lists]
  # Create a heap with a slot for each list.
  heap = range(len(Lists))
  for i in heap:
    # Heap elements are (key, value) pairs of (element, List index)
    heap[i] = (Lists[i][0], i)
  BuildMinHeap(heap)
  # Last tells the index of the list with the min element.
  last = heap[0][1]
  #############################################
  # O(n)
  #############################################
  merged = range(sum(sizes))
  for i in merged:
    # Next-in-order element at root of min heap.
    merged[i] = heap[0][0]
    while sizes[last] == 0:
      last = (last+1) % len(Lists)
    # Fill the hole in the heap and Min-Heapify.
    heap[0] = (Lists[last][-sizes[last]], last)
    sizes[last] -= 1
    #############################################
    # O(log k) where k = len(Lists)
    #############################################
    MinHeapify(heap)
  return merged

Leave a Reply