Merge K sorted arrays
Merge 'k' sorted linked lists and return it as one sorted list. Analyze and describe its complexity.
# Definition for singly-linked list.
# class ListNode(object):
# def __init__(self, x):
# self.val = x
# self.next = None
'''
heap[k] <= heap[2*k+1] and heap[k] <= heap[2*k+2] for all k, counting elements from zero. For the sake of comparison, non-existing elements are considered to be infinite. The interesting property of a heap is that its smallest element is always the root, heap[0].
'''
class Solution(object):
def mergeKLists(self, lists):
"""
:type lists: List[ListNode]
:rtype: ListNode
"""
# import heapq module
import heapq
h = [] # this list will define the heap using the heapq module in Python
k = len(lists) # number of k sorted lists
mergedList = [] # return this
#================================================================================
# base case
#================================================================================
if not k:
return mergedList
#================================================================================
# Store the pointers to the next Node in an iterable list
# This will help us to quickly come to the next Node in the correct kth list, after the smallest
# element is popped from the heap.
#================================================================================
iterList = [None for i in xrange(k)]
#================================================================================
# Initially: Just buld a heap with the first elements from all k sorted lists
# Each heap element is a tuple of (element value, listIndex)
# The listIndex in the tuple will help us identify the which list number
# out of those k lists to go to after this tuple gets popped.
#================================================================================
for listIndex in xrange(k):
list = lists[listIndex]
if list:
tup = tuple((list.val, listIndex))
#print 'tup(start): ', tup
heapq.heappush(h, tup)
iterList[listIndex] = list.next # store the next pointer
#================================================================================
# Main task: 1. Pop the smallest tuple from the heap.
# 2. Find the listIndex from that popped tuple
# 3. Find the list Node from the pointer stored in the iterList for that listIndex
# 4. Push that list node on to the heap
# 5. The heapq will maintain the smallest element always at the top
#================================================================================
while h:
elem = heapq.heappop(h)
listIndex = elem[1]
#print 'min elem: ', elem[0], ' listIndex = ', listIndex
mergedList.append(elem[0])
if iterList[listIndex] is not None:
tup = tuple((iterList[listIndex].val, listIndex))
#print 'tup: ', tup
heapq.heappush(h, tup)
iterList[listIndex] = iterList[listIndex].next
#print 'merged List: ', mergedList
return mergedList
Last updated
Was this helpful?