Merge k-sorted lists - Medium

Problem - Given k-sorted lists, merge them into a single sorted list.

A daft way of doing this would be to copy all the list into a new array and sorting the new array. ie O(n log(n))

The naive method would be to simply perform k-way merge similar to the auxiliary method in Merge Sort. But that is reduces the problem to a minimum selection from a list of k-elements. The Complexity of this algorithm is an abysmal O(nk).

Here's how it looks in Python. We maintain an additional array called Index[1..k] to maintain the head of each list.

	def merge(arr):
	k = len(arr)
	index = [0] * k
	lst = []
	while True:
	mni = None
	for i in xrange(k):
	if index[i] < len(arr[i]) and (mni == None or arr[i][index[i]] < arr[mni][index[mni]]):
	mni = i
	if mni == None: break
	lst.append(arr[mni][index[mni]])
	index[mni] += 1
	return lst

view raw gistfile1.txt hosted with ❤ by GitHub

We improve upon this by optimizing the minimum selection process by using a familiar data structure, the Heap! Using a MinHeap, we extract the minimum element from a list and then push the next element from the same list into the heap, until all the list get exhausted.

This reduces the Time complexity to O(nlogk) since for each element we perform O(logk) operations on the heap. An important implementation detail is we need to keep track of the origins of the elements of the heap, ie the list it came from and it's index. Therefore, I use a 3 member tuple (item, list-index, item-index) and push it onto the heap.

Fortunately for us, Python has the heapq module for implementing a heap. So here's the code.

	def merge2(arr):
	heap = [(l[0],i,0) for i,l in enumerate(arr) if len(l) > 0]
	heapq.heapify(heap)
	lst = []
	while heap:
	item,lst_index,item_index = heapq.heappop(heap)
	lst.append(item)
	if item_index + 1 < len(arr[lst_index]):
	heapq.heappush(heap,(arr[lst_index][item_index+1],lst_index,item_index+1))
	return lst

view raw gistfile1.txt hosted with ❤ by GitHub

The entire code can be found here.