"""
https://leetcode.com/problems/merge-k-sorted-lists/
Merge k sorted linked lists and return it as one sorted list. Analyze and describe its complexity.
Example:
Input:
[
1->4->5,
1->3->4,
2->6
]
Output: 1->1->2->3->4->4->5->6
"""
from heapq import heappop, heapify
# Definition for singly-linked list.
class ListNode:
def __init__(self, x, next=None): # 注意我这里next参数
self.val = x
self.next = next
class Solution:
def mergeKLists(self, lists):
"""
:type lists: List[ListNode]
:rtype: ListNode
"""
h = []
for node in lists:
while node:
h.append(node.val)
node = node.next
if not h: # 注意这里的边界值 [[]]
return None
heapify(h) # inplace 建堆
root = ListNode(heappop(h))
curnode = root
while h:
nextnode = ListNode(heappop(h))
curnode.next = nextnode
curnode = nextnode
return root
def to_list(head):
"""list to []"""
res = []
curnode = head
while curnode:
res.append(curnode.val)
curnode = curnode.next
return res
def test_merge():
l1 = ListNode(1, ListNode(4, ListNode(5))) # 1->4->5
l2 = ListNode(1, ListNode(3, ListNode(4))) # 1->3->4
s = Solution()
lists = [l1, l2]
head = s.mergeKLists(lists)
print(to_list(head))
if __name__ == '__main__':
test_merge()
