23. Merge K sorted Lists

Problem Link

Tutorial

Use priority queue to store the first element of each list. Then every time in a loop, pop one and push in the next one in the list.

In python, priority queue is defaultly to pop the smallest element. If you want to pop the largest, just insert the negative value. If the comparing is complicated, just push in ((keys in order), elem), so that we not only have the entire element, but also the keys to compare with in a tuple.

Solution

# Definition for singly-linked list.
# class ListNode(object):
#     def __init__(self, x):
#         self.val = x
#         self.next = None
from Queue import PriorityQueue
class Solution(object):
    def mergeKLists(self, lists):
        """
        :type lists: List[ListNode]
        :rtype: ListNode
        """
        q = PriorityQueue()
        for node in lists:
            if node:
                q.put((node.val, node))
        ret = p = ListNode(None)
        while q.qsize() > 0:
            (val, node) = q.get()
            if node.next:
                q.put((node.next.val, node.next))
            p.next = node
            p = p.next
        return ret.next