greentea31's Note

https://leetcode.com/problems/merge-k-sorted-lists/

풀이 과정

k개의 오름차순으로 정렬되어 있는 링크드 리스트를 입력받아 1개의 오름차순 정렬 링크드 리스트를 반환하는 문제이다.

우선순위 큐를 이용하여 문제를 해결하였다. heapq 모듈을 사용하였는데 heapq는 최소 힙을 구현한 모듈이다. (노드 값, 인덱스, 링크드 리스트) 튜플을 최소 힙에 다 넣고 최소 힙에서 나오는 순서대로 연결 순서를 조절해 줌으로써 하나의 링크드 리스트로 통합하였다.

튜플에서 인덱스 값은 필요가 없는 값인데, 튜플 (a, b, c), (d, e, f)를 heap에 heappush시에 a와 d를 비교해서 작은 값 -> 같으면 b와 e를 비교해서 작은 값 -> 같으면 c와 f를 비교해서 작은 값이 힙에서 위로 가게끔 한다. 그런데 인덱스 값을 쓰지 않으면 노드 값이 같은 두 링크드 리스트의 튜플 (노드 값 1, 링크드 리스트 1), (노드 값 2, 링크드 리스트 2)을 비교할 때 노드 값 1 == 노드 값 2 이므로 링크드 리스트 1, 링크드 리스트 2를 비교하게 하는데 링크드 리스트 1 < 링크드 리스트 2 같은 연산은 정의되어 있지 않아서 오류가 발생한다.

그래서 (노드 값 1, 인덱스 값 1, 링크드 리스트 1), (노드 값 2, 인덱스 값 2, 링크드 리스트 2) 같이 lists 내의 인덱스 값을 튜플에 추가해서 노드 값이 같으면 lists에서 먼저 나온 링크드 리스트가 최소 힙의 루트 쪽에 가깝게 배열되게 설정한 것이다.

소스 코드

class Solution:
    def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
        root = res = 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]
            res.next = node[2]
            res = res.next
            
            if res.next:
                heapq.heappush(heap, (res.next.val, idx, res.next))
        
        return root.next