[Leetcode 아주 상세한 문제풀이] 23. Merge k Sorted Lists – 코드 line 별 설명

Leetcode 문제

이번에 풀어볼 문제는 리트코드 23번 문제 Merge k Sorted Lists 다.

우선 문제를 살펴보자.

 

리트코드 23번 문제 Merge k Sorted Lists 링크

Merge k Sorted Lists - LeetCode

Merge k Sorted Lists - LeetCode

 

You are given an array of k linked-lists lists, each linked-list is sorted in ascending order.

Merge all the linked-lists into one sorted linked-list and return it.

Example 1:
Input: lists = [[1,4,5],[1,3,4],[2,6]]
Output: [1,1,2,3,4,4,5,6]
Explanation: The linked-lists are:
[
  1->4->5,
  1->3->4,
  2->6
]
merging them into one sorted list:
1->1->2->3->4->4->5->6

Example 2:
Input: lists = []
Output: []

Example 3:
Input: lists = [[]]
Output: []
 
Constraints:
k == lists.length
0 <= k <= 104
0 <= lists[i].length <= 500
-104 <= lists[i][j] <= 104
lists[i] is sorted in ascending order.
The sum of lists[i].length will not exceed 104.

 

 

문제 설명

문제를 간략히 살펴보자. 인풋으로 linked list들로 이루어진 배열이 주어진다. 이때 배열 내 각 linked list들은 오름차순으로 정렬되어 있다. 문제에서 요구하는 바는 인풋으로 주어진 linked list들을 합쳐 오름차순으로 나타내라는 것이다.

문제 예시 1번을 살펴보면 인풋으로 3 개의 linked list [1, 4, 5], [1, 3, 4], [2, 6]이 주어졌다. 문제의 조건에 따라 각 linked list는 오름차순으로 배열되어있다. 이때 이 3개의 linked list를 하나로 합쳐 오름차순으로 나타내야 한다. 이 경우 결과는 [1, 1, 2, 3, 4, 4, 5, 6]으로 얻어진다. 단순 배열이라면 주어진 배열을 합쳐 오름차순 배열하는 것이 어렵지 않지만 이 경우 linked list들을 합쳐 또 하나의 linked list를 만들되 오름차순이 되어야 한다는 것을 명심하자. 이제 문제를 풀어보도록 하자.

 

 

문제 접근법

이 문제를 해결하기 위해서 min heap을 이용할 것이다.

이때 heap은 linked list들의 노드를 담아두는데 heap의 성질에 따라 노드 크기 순대로 배열이 된다.

좀 더 정확히는 가장 작은 값을 가진 노드가 가장 상단에 위치하게 된다.

우리는 각 linked list에 대하여 head (첫번째 노드)를 heap에 담는다.

각 linked list에서 가장 작은 값 (이미 오름차순으로 배열되어 있으므로)들이 heap에 담겼다.

그리고 heap에서 가장 작은 element를 제거해 merged linked list에 더한다.

그 후 제거된 element가 인풋으로 주어졌던 linked list에서 다음 노드가 있다면 그 노드를 heap에 담는다.

이 과정을 heap에서 모든 노드가 제거될 때 까지 반복한다.

예를 들어 문제 예시 1번과 같이 [[1, 4, 5], [1, 3, 4], [2, 6]]이 주어졌다고 생각해보자.

위에서 설명한대로 각 linked list에서 가장 작은 element를 heap 담는다. [1, 1, 2]

Heap의 성질에 따라 가장 작은 값이 가장 위에 위치한다.

Pop을 이용하여 가장 작은 1을 제거해 merged linked list에 담으면 heap에는 [1, 2]이 남는다.

그리고 제거한 1 뒤에 있던 4 (1->4->5)를 heap에 담는다. heap=[1, 2, 4]

다음으로 heap에서 가장 작은 1을 제거하여 merged linked list에 담으면 heap에는 [2, 4]가 남는다.

마찬가지로 제거한 1 뒤에 있던 3 (1->3->4)를 heap에 담는다. heap=[2, 3, 4]

이 과정을 heap에 아무것도 남아있지 않을 때까지 반복한다.

와중에 heap에서 가장 작은 element를 merged linked list에 '오름차순'으로 넣었으므로 오름차순 배열된 merged linked list를 얻을 수 있다.

이제 코드로 한 번 살펴보자.

 

 

문제 풀이

Heap을 이용하기 위하여 heapq module를 가져온다. (line 1)

Linked list를 표현하기 위한 클래스를 선언해준다. (line 3-10)

이는 파이썬에서 일반적으로 사용되는 표현식이니 익숙해 지도록 하자.

Linked list로 이루어진 배열을 입력받기 위한 함수를 정의해 준다. (line 12)

Heap을 이용할 계획이므로 heap을 만들어 준다. (line 13)

인풋으로 입력받은 각 linked list를 for loop를 이용하여 탐색한다. (line 14)

만약 linked list 내 노드가 존재할 때 head값을 heap에 넣어준다. (line 15-16)

앞서 설명한대로 각 linked list의 가장 작은 값이 heap에 들어간다.

이제 각 linked list 들을 합쳐 merged linked list를 만들기 위한 작업을 위해 dummy node를 하나 만들어 현재 위치를 그 dummy node에 위치시킨다. (line 18-19)

위에서 설명한대로 heap 안에 노드값이 존재하는한 pop으로 가장 위에 위치한 값 (=가장 작은 값)을 제거한다. (line 22)

그리고 그 값을 merged linked list의 처음 (current.next, dummy 바로 다음)에 넣어준다. (line 23)

그리고 다음 값을 넣기 위하여 current 위치를 다음으로 옮긴다. (line 24)

만약 pop으로 제거한 값에 연결된 다음 노드가 있다면 그것을 heap에 넣어준다. (line 26-27)

이 모든 과정을 거치고 나면 heap에는 아무것도 남아있지 않다.

그리고 heap에서 제거된 값들이 merged linked list에 오름차순으로 배열되어 있다.

따라서 우리는 오름차순으로 병합된 linked list를 출력하고자 하므로 dummy의 다음값 (dummy는 merged linked list head 앞에 의도적으로 만들어 놓은 것이므로)을 출력하고 종료한다.

 

 

문제 풀이 코드

import heapq
 
class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next
 
    def __lt__(self, other):
        # This ensures the ListNode can be compared based on its value
        return self.val < other.val
 
def mergeKLists(lists):
    heap = []
    for l in lists:
        if l:
            heapq.heappush(heap, l)
 
    dummy = ListNode()
    current = dummy
 
    while heap:
        node = heapq.heappop(heap)
        current.next = node
        current = current.next
 
        if node.next:
            heapq.heappush(heap, node.next)
 
    return dummy.next
 
def list_to_linkedlist(arr):
    # Helper function to convert a list to a linked list
    dummy = ListNode()
    current = dummy
    for val in arr:
        current.next = ListNode(val)
        current = current.next
    return dummy.next
 
def linkedlist_to_list(node):
    # Helper function to convert a linked list to a list
    arr = []
    while node:
        arr.append(node.val)
        node = node.next
    return arr
 
# Test the function with the given examples
lists1 = [list_to_linkedlist(l) for l in [[1,4,5], [1,3,4], [2,6]]]
lists2 = []
lists3 = [list_to_linkedlist(l) for l in [[]]]
 
merged1 = mergeKLists(lists1)
merged2 = mergeKLists(lists2)
merged3 = mergeKLists(lists3)
 
print("Merged list 1:", linkedlist_to_list(merged1))
print("Merged list 2:", linkedlist_to_list(merged2))
print("Merged list 3:", linkedlist_to_list(merged3))