[Leetcode] 리트코드 - k개 정렬 리스트 병합(merge-k-sorted-lists) 파이썬(python) 풀이

k개 정렬 리스트 병합

문제

k개의 정렬된 리스트를 1개의 정렬 리스트로 병합하라

입력

[1 → 4 → 5, 1 → 3 → 4, 2 → 6]

출력

1 → 1 → 2 → 3 → 4 → 4 → 5 → 6

코드

from typing import List
import heapq

class ListNode:
    def __init__(self, x):
        self.val = x
        self.next = None

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

l1 = ListNode(1)
l1.next = ListNode(4)
l1.next.next = ListNode(5)

l2 = ListNode(1)
l2.next = ListNode(3)
l2.next.next = ListNode(4)

l3 = ListNode(2)
l3.next = ListNode(6)

Solution().mergeKLists([l1, l2, l3])

풀이

for i in range(len(lists)):
    if lists[i]:
        heapq.heappush(heap, (lists[i].val, i, lists[i]))

각 연결 리스트의 루트를 힙에 저장한다. 파이썬의 heapq 모듈은 최소 힙을 지원하며, 따라서 list.val이 작은 순서대로 pop()할 수 있다. 그런데 이렇게 저장하면 TypeError: '<' not supported between instances of 'ListNode' and 'ListNode'라는 에러가 발생한다. '중복된 값을 넣을 수 없다'라는 뜻이다.

이 문제의 예제로 제시한 입력값은 3개의 연결 리스트 중 첫 번째와 두 번째의 루트가 각각 1로 동일하다. 이렇게 동일한 값은 heappush() 함수에서 에러를 발생하기 때문에 중복된 값을 구분할 수 있는 추가 인자가 필요하다. 따라서 heappush의 두번째 인자에 처음에는 value를 넣고 두번째로 index를 삽입하여 에러가 나지 않게 해준다.

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))

result.next = node[2]를 통해 힙에서 빼온 node를 result의 다음 노드가 가르킬 수 있게 해준다.

출처

파이썬 알고리즘 인터뷰 (글 : 박상길 그림 : 정진호) [책만]