[파이썬 알고리즘 인터뷰] 27. k개 정렬 리스트 병합

 

 이 포스팅은 파이썬 알고리즘 인터뷰를 공부하면서 정리한 내용입니다.

 

 출처 : 파이썬 알고리즘 인터뷰

 코드 출처 : 파이썬 알고리즘 인터뷰 깃허브

 

---

27. k개 정렬 리스트 병합

leetcode 23. Merge K Sorted Lists 문제입니다.

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

 

 

풀이

 우선순위 큐를 활용해서, 최소값에 해당하는 node.val을 추출하고, next가 존재하면 next.val를 우선순위 큐에 다시 넣어줍니다. 파이썬에서 우선순위 큐는 heapq 모듈을 사용합니다.

 

def mergeKLists(self, lists):
    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