[leetcode / 23] Merge k Sorted Lists

in Problem Solving on leetcode


날짜: 2021년 8월 28일
카테고리: 우선순위 큐
태그: Hard, 23, 파이썬

leetcode 23 - Merge k Sorted Lists

입출력 예시

예제 입력예제 출력
lists = [[1,4,5],[1,3,4],[2,6]][1,1,2,3,4,4,5,6]
lists = [][]
lists = [[]][]

코드

import heapq

# Definition for singly-linked list.
# class ListNode:
#     def __init__(self, val=0, next=None):
#         self.val = val
#         self.next = next
class Solution:
    def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
        root = answer = ListNode(None)
        heap = list()
        
        for i in range(len(lists)):
            if lists[i]:
                heapq.heappush(heap, (lists[i].val, i ,lists[i]))
        
        while heap:
            node = heapq.heappop(heap)
            index = node[1]
            answer.next = node[2]
            
            answer = answer.next
            if answer.next:
                heapq.heappush(heap, (answer.next.val, index, answer.next))
                
        return root.next
파이썬 알고리즘 인터뷰 10-2

풀이 과정

root와 answer이라는 연결 리스트와 heap이라는 리스트를 하나 생성

입력 받은 lists의 길이만큼 반복하여 heapq.heappush() 함수를 활용하여 heap에 노드 정보를 push
(중복되는 값에 대한 에러를 방지하기 위한 형태)

이제 heapq.heappop()으로 pop하면 가장 작은 값의 노드부터 차례로 연결 리스트 answer에 차례로 노드 정보를 push
힙에 아무 값도 남지 않을 때까지 반복하여 root.next를 반환

반성

  • 내가 이 문제를 보고 바로 우선순위 큐를 활용하기 위해 heapq라이브러리를 import할 생각이 들 수 있을까…


© 2021. by hminkim