2020年5月8日 星期五

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