Lintcode 104. Merge K Sorted Lists (Medium) (Python) (有问题)

Merge K Sorted Lists

Description:

Merge k sorted linked lists and return it as one sorted list.

Analyze and describe its complexity.

Example
Given lists:

[
2->4->null,
null,
-1->null
],
return -1->2->4->null.

Code: 50%时未通过

"""
Definition of ListNode
class ListNode(object):

    def __init__(self, val, next=None):
        self.val = val
        self.next = next
"""
class Solution:
    """
    @param lists: a list of ListNode
    @return: The head of one sorted list.
    """
    def mergeKLists(self, lists):
        # write your code here
        import heapq
        heap = []
        for i in lists:
            if i != None:
                heap.append((i.val, i))
        heapq.heapify(heap)
        res = ListNode(0)
        cur = res
        while heap:
            val, node = heapq.heappop(heap)
            cur.next = node
            cur = cur.next
            if node.next:
                print(node.val)
                heapq.heappush(heap, (node.next.val, node.next))
        return res.next