[파이썬 알고리즘 인터뷰][데크/우선순위큐] k개 정렬 리스트 병합

이 글은 "파이썬 알고리즘 인터뷰 (박상길 지음)"을 읽고 주관적으로 요약한 글입니다. 

---

출처 : https://www.onlybook.co.kr/entry/algorithm-interview

 

문제 정의

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

 

책에서 구현된 코드

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는 사용시 list를 인자로 받음
  • import시 collection, itertools 등등에 속해있지 않고 바로 import
  • 인자로 tuple을 받음

 

내가 구현한 코드

def mergeKLists(lists: list[ListNode]) -> ListNode:
    pq = []
    res = root = ListNode(0)

    for node in lists:
        while node:
            heapq.heappush(pq, node.val)
            node = node.next

    while pq:
        res.next = ListNode(heapq.heappop(pq))
        res = res.next

    return root.next

• 내가 구현한 코드

• 책에서 구현한 코드

• heapq.heappush
  
  • 책에서 같은 값으로 push가 안된다고 되어있는데, 중복된 값도 push가 된다