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

 

<--- 23. Merge k Sorted Lists --->

# 문제: k개의 정렬된 리스트를 1개의 정렬된 리스트로 병합하라.

 

입력

[   1->4->5,   1->3->4,   2->6 ]

 

출력

1->1->2->3->4->4->5->6

 

 

# 풀이

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

import heapq
 
class ListNode(object):
  def __init__(self, val=0, next=None):
    self.val = val
    self.next = None
    
def merge_lists(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