Leetcode題解 Python:Merge k Sorted Lists

合併 k 個 linked list 成一個 linked list,並且值由小到大排序。
說到要保持由小而大,不外乎就想到 min heap,關於 heap 的介紹,這裡就不多說了。
因為每次新的加入就得重新找出最小的,選擇 heap 是快多了。
不過不論是 heap 還是 sort() 都無法對於 ListNode 做比較,所以不是重新生成ListNode,
不然就是得找另外一種的保存方式。
我選擇用 dict() 以Node的值為鍵去存放 ListNode。但可能會出現同值有二個以上的ListNode,所以得用list存放,
同值的順序並不重要,抓哪個都可以。為了方便,使用 collections.defaultdict(list)。
用 heap 的 list 只有放 ListNode.val,只要拿值去對就可以取出下一個要接的 ListNode。

Python

class Solution:
def mergeKLists(self, lists: List[ListNode]) -> ListNode:
# 使用 min heap
import heapq
import collections
nodeMap = collections.defaultdict(list)
stack = []
for node in lists:
if node:
nodeMap[node.val].append(node)
heapq.heappush(stack, node.val)
if not stack: return None
root = nowNode = ListNode()
while stack:
nowNode.next = nodeMap[heapq.heappop(stack)].pop()
nowNode, nextNode = nowNode.next, nowNode.next.next
if nextNode:
nodeMap[nextNode.val].append(nextNode)
heapq.heappush(stack, nextNode.val)
return root.next