10-line python solution with priority queue

  • 49
    from Queue import PriorityQueue
    class Solution(object):
        def mergeKLists(self, lists):
            dummy = ListNode(None)
            curr = dummy
            q = PriorityQueue()
            for node in lists:
                if node: q.put((node.val,node))
            while q.qsize()>0:
                curr.next = q.get()[1]
                if curr.next: q.put((curr.next.val, curr.next))
            return dummy.next

  • 13

    Whaaaaat. How did I not know the Queue module.

  • 0

    it may not be better than heapq, Queue does not has a peek() function...

  • 0

    This is really good. Thanks for sharing!

  • 4

    It's better to use heapq in my opinion. The Queue class in python is implemented for synchronized tasks, not as a data structure only. From the documentation :the Queue module implements multi-producer, multi-consumer queues. It is especially useful in threaded programming when information must be exchanged safely between multiple threads.

  • 0

    So would the complexity of this solution be O(Nk logk) where N is the number of lists and k is the average size of a list?

  • 0

    Thank you for sharing this concise answer.
    May I ask a question, how you learn about PriorityQueue?
    It might be not a rare data structure but it's also not as popular as the basic ones like dict, heap.

  • 0

    @StefanPochmann said in 10-line python solution with priority queue:

    Whaaaaat. How did I not know the Queue module.

    Haha, so did I, professional Noob.

  • 0

    @cbmbbz In the event that two or more of the lists have the same val, this code will error out since the queue module will compare the second element in the priority queue which is a ListNode object (and this is not a comparable type).

    To solve for this issue, I have stored (node.val, list_idx, node) to account for this edge case.

    from queue import PriorityQueue
    # Definition for singly-linked list.
    # class ListNode:
    #     def __init__(self, x):
    #         self.val = x
    #         self.next = None
    class Solution:
        def mergeKLists(self, lists):
            k = len(lists)
            q = PriorityQueue(maxsize=k)
            dummy = ListNode(None)
            curr = dummy
            for list_idx, node in enumerate(lists):
                if node: q.put((node.val, list_idx, node))
            while q.qsize() > 0:
                poped = q.get()
                curr.next, list_idx = poped[2], poped[1]
                curr = curr.next
                if curr.next: q.put((curr.next.val, list_idx, curr.next))
            return dummy.next

Log in to reply

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