[Leetcode] 23. Merge k Sorted Lists

You are given an array of k linked-lists lists, each linked-list is sorted in ascending order.

Merge all the linked-lists into one sorted linked-list and return it.

---

 

Example 1:

Input: lists = [[1,4,5], [1,3,4], [2,6]] Output: [1,1,2,3,4,4,5,6] Explanation: The linked-lists are: [ 1->4->5, 1->3->4, 2->6 ] merging them into one sorted list: 1->1->2->3->4->4->5->6

Example 2:

Input: lists = [] Output: []

Example 3:

Input: lists = [[]] Output: []

---

Constraints:

• k == lists.length
• 0 <= k <= 10^4
• 0 <= lists[i].length <= 500
• -10^4 <= lists[i][j] <= 10^4
• lists [i]is sorted in ascending order.
• The sum of lists[i].length won't exceed 10^4.

---

풀이

class Solution:
    def mergeKLists(self, lists: List[ListNode]) -> ListNode:
        root = result = ListNode(None)
        heap = []
        # 각 연결리스트의 루트를 힙에 저장
        for i in range(len(lists)):
            if lists[i]:
                heapq.heappush(heap,(lists[i].val,i,lists[i]))
        # 힙 추출 이후 다음 노드는 다시 저장
        while heap:
            node = heapq.heappop(heap)
            idx = node[1]
            result.next = node[2]
            
            result = result.next
            if result.next:
                heapq.heappush(heap,(result.next.val , idx, result.next))
        return root.next

결과

---

문제 풀이

이 문제는 우선순위 큐를 사용해 쉽게 풀 수 있는 문제이다. 이 문제의 예제로 제시한 입력값은 3개의 연결 리스트 중 첫 번째와 두 번째의 루트가 각 1로 동일하다. 이렇게 동일한 값은 heappush() 함수에서 에러가 발생하기 때문에 중복된 값을 구분할 수 있는 추가 인자가 필요하다.

heapq.heappush(heap,(lists[i].val,i,lists[i]))

 

이제 heappop()으로 값을 추출하면 가장 작은 노드의 연결 리스트부터 차례대로 나오게 되며, 이 값을 결과가 될 노드 result에 하나씩 추가한다. k개의 연결 리스트가 모두 힙에 계속 들어 있어야 그중에서 가장 작은 노드가 항상 차례대로 나올 수 있으므로 추출한 연결 리스트의 그다음 노드를 다시 힙에 추가한다.

heapq.heappush(heap,(result.next.val , idx, result.next))