Python solution using heapq (priority queue)


  • 0
    B
    # 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
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.