Python 알고리즘 인터뷰 - 문제27. k개 정렬 리스트 병합

leetcode 23. Merge k Sorted Lists - https://leetcode.com/problems/merge-k-sorted-lists/description/

 

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

 

예제

입력

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

 

출력

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

 

풀이

우선순위 큐를 사용하면 문제를 간단히 해결할 수 있다.

python에서는 우선순위 큐 문제 해결을 위해 headq를 사용하므로 해당 모듈을 사용하여 구현하도록 한다.

headq에 각 연결리스트를 추가한다. 값은 tuple로 (node.val, node)로 전달하여 node.val을 기준으로 정렬하도록 한다.

for lst in lists:
	heapq.heappush(heap, (lst.val, lst))

 

해당 코드는 다음과 같이 heap에서 val이 동일한 데이터가 존재할 경우 에러가 발생한다. tuple을 비교할 시 첫번째 데이터가 동일할 때 두번째 데이터를 비교하는데 Node와 Node를 비교하는 함수가 구현되어 있지 않기 때문이다. 

따라서 다음과 같이 값이 동일할 때 비교할 수 있도록 index값을 추가하여 에러를 방지한다.

 

for i in range(len(lists)):
	heapq.heappush(heap, (lists[i].val, i, lists[i]))

이후 heap에서 값을 pop한 후 result가 해당 노드를 계속 연결한 후 이동한다.

result는 다음 노드를 튜블로 만들어서 heap에 넣는데, heap은 result.next.val의 우선순위에 데이터를 배치한다. 따라서 힙에는 항상 노드의 값이 가장 작은 순으로 pop이 된다.

동일하게 위의 내용을 진행하면 다음과 같이 result는 값이 작은 순서대로 나열되고, heap에는 데이터가 하나씩 사라진다.

 

이를 코드로 나타내면 다음과 같다.

result = ListNode(None)
while heap:
	node = heapq.heappop(heap)
    idx = node[1]
    result.next = node[2]
    
    result = result.next
    if result.next:
    	headq.heappush(heap, (result.next.val, idx, result.next)

 

전체 코드는 다음과 같다.

def mergeKLists(lists: list[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