설명

LeetCode의 Merge K Sorted Lists 문제다.

말 그대로 K 개의 정렬된 리스트를 병합하는 문제인데 리스트는 연결 리스트 형태로 주어지며 오름차순으로 정렬되어 있다.

예를 들어 [1,2,3], [4,5,6], [2,3,4]가 주어진다면 [1,2,2,3,3,4,4,5,6]으로 정렬되어야 한다는 것이다. 언뜻 보기에는 간단한 문제 같지만 문제의 요구사항을 살펴보면 이 정렬된 리스트는 10000개까지 주어질 수 있으며 각 리스트는 500개까지 원소를 포함할 수 있다. 즉 최악의 경우 5000000개의 원소를 정렬해야 하는 것이다.

풀이

단순 반복 비교(3848 ms)

첫번째로 생각한 풀이는 시간 복잡도에 신경쓰지 않고 일단 푸는것에 집중했다. 하지만 너무 신경을 쓰지 않았던 것인지 3848 ms라는, 통계에도 잡히지 않는 최악의 성능을 얻게 되었다.

이 풀이의 로직은 주어진 리스트 중 하나를 기준으로 잡아서 나머지 리스트를 오름차순으로 병합시키는 것이다.

class Solution:
    def mergeKLists(self, lists: List[ListNode]) -> ListNode:
        # 단순 예외 처리
        if not lists:
            return None
        if len(lists) < 2:
            return lists[0]
   
        # 병합될 기준 리스트 탐색(빈 리스트가 주어질 수 있기 때문)
        start = None
        idx = 0
        for index, listNode in enumerate(lists):
            if listNode is not None:
                start = lists[index]
                idx = index
                break
    
        # 병합할 리스트들을 탐색
        for listNode in lists[index+1:]:
            # 빈 리스트는 무시
            if not listNode:
                continue
            
            # 병합할 리스트의 요소들을 전부 탐색
            current = listNode # 병합할 리스트의 시작 요소
            while current:
                prev = None
                compare = start # 기준 리스트의 요소
                isMerged = False
                # 기준 리스트의 요소와 비교하여 삽입
                while compare:
                    if current.val <= compare.val:
                        isMerged = True
                        insertedNode = ListNode(current.val, compare)
                        # 첫 번째 요소일 때는 해당 요소 이전에 삽입
                        if prev is None:
                            start = insertedNode
                        # 만약 두 번째 요소 또는 그 이후라면 이전 노드와 연결
                        else:
                            prev.next = insertedNode
                        break
                    else:
                        prev = compare
                    
                    # 기준 리스트의 다음 요소와 비교해야 하므로 이동
                    compare = compare.next
                
                # 삽입할 요소가 기준 리스트의 모든 원소보다 클 경우 맨 뒤에 삽입
                if not isMerged and prev:
                    prev.next = ListNode(val=current.val)
                
                current = current.next
            
        return start

분명 코드도 길고 로직도 복잡해서 좋지 않은 풀이인 것은 분명하다. 그럼 어떤 방법을 시도해 볼 수 있을까?

리스트의 정렬 메서드 활용(96 ms)

파이썬 리스트의 장점 중 하나는 자체적으로 정렬 메서드를 지원하고 있고 시간 복잡도도 O(N*logN)에 불과하다는 것이다. 이를 어떻게 활용할 수 있을까? 이 문제에서 연결 리스트의 노드를 나타내는 Node 클래스가 현재 노드의 값을 의미하는 val 필드와 다음 노드를 가리키는 next 필드만 포함하는 간단한 구조라는 것을 활용할 수 있다.

그래서 모든 연결 리스트의 노드들의 값을 하나의 리스트에 담아서 정렬한 후 다시 그대로 노드를 생성하면 결과적으로 모든 노드들이 오름차순으로 정렬된 연결 리스트를 구축할 수 있을 것이다.

class Solution:
    def mergeKLists(self, lists: List[ListNode]) -> ListNode:
        # 모든 노드들의 값을 담을 리스트 생성
        nodes = []
        
        # 모든 노드들의 값을 리스트에 삽입
        for listNode in lists:
            while listNode:
                nodes.append(listNode.val)
                listNode = listNode.next
            
        # 정렬된 값들로 다시 노드를 생성, 하나의 연결 리스트 구축    
        start = current = ListNode()
        for node in sorted(nodes):
            current.next = ListNode(val=node)
            current = current.next
        
        return start.next

우선순위 큐를 이용한 정렬(104 ms)

위의 풀이와 비슷한 방법이지만 이번에는 Python3의 heapq 모듈을 사용해서 우선순위 큐를 활용했다.

import heapq


class Solution:
    def mergeKLists(self, lists: List[ListNode]) -> ListNode:
        # 노드들의 값을 담을 큐 생성
        queue = []
        
        # 모든 노드의 값을 큐에 삽입
        for listNode in lists:
            while listNode:
                heapq.heappush(queue, listNode.val)
                listNode = listNode.next
                
        # 우선순위 큐에서 요소를 하나씩 빼서 연결 리스트 구축        
        start = current = ListNode()
        for _ in range(len(queue)):
            current.next = ListNode(val=heapq.heappop(queue))
            current = current.next
            
        return start.next

heapq 모듈의 heappop 메서드는 우선순위 큐에서 가장 작은 값을 반환한다. 이는 고정되어 있으며 변경할 수 없다. 그렇기 때문에 역순으로 활용하고자 한다면 튜플을 이용해서 내림차순의 값을 활용해야 할 것이다.

후기

위처럼 노드의 값들만 한 곳에 모아두었다가 나중에 이를 기준으로 다시 노드를 생성하는 방법은 빠르고 간단해서 좋은 풀이같지만 연결 리스트의 관점에서 본다면 좀 애매한 것 같다.

정렬된 배열을 병합하는 문제였다면 괜찮겠지만 지금은 각 노드들이 연결된 연결 리스트를 병합해야 하는데 만약 노드가 간단한 정수값만 가진 게 아니라 매우 복잡하고 큰 필드도 포함하고 있다면 어떨까? 그때는 지금처럼 단순히 값을 복사해서 리스트에 저장하고 정렬시키는 것만으로 끝나지 않을 것이다.

참고자료

• 파이썬 알고리즘 인터뷰