[LeetCode] #23. Merge k Sorted Lists (python)

문제

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

Merge k Sorted Lists - LeetCode

k개의 정렬된 연결리스트들을 오름차순으로 정렬된 하나의 연결리스트로 합쳐야한다.

 

예시 입출력

Input:
lists = [[1,4,5],[1,3,4],[2,6]]
Output:
[1,1,2,3,4,4,5,6]

 

 

풀이

 

책 풀이

우선순위 큐를 이용한 리스트 병합

import heapq
from typing import List


# Definition for singly-linked list.
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

책에서는 우선순위 큐를 사용해서 풀었다. 우선순위 큐를 사용해서 풀기 위해 heapq 모듈을 사용한다. heapq 모듈은 최소 힙을 지원한다. 

 

우선순위 큐란?

특정 조건에 따라 우선 순위가 가장 높은 요소가 추출되는 자료형이다. 

 

heapq의 기본 함수들

heapq.heappush(heap, item) : heap에 item 추가

heapq.heappop(heap) : heap에서 가장 작은 원소를 pop

 

 

각 연결리스트의 루트를 힙에 저장한다.

저장한 힙에서 heappop을 하면 가장 작은 노드의 연결리스트가 추출된다.

해당 값을 result 노드에 하나씩 연결하고 나머지 노드들은 다시 힙에 추가한다.

최종 결과인 root.next를 반환한다. (root는 None으로 시작하기 때문에 root.next를 반환)

 

참고한 책 : 파이썬 알고리즘 인터뷰