Python으로 정렬된 K개의 리스트를 병합하는 방법 – 힙(Heap) 활용 가이드

Python으로 정렬된 K개의 리스트 병합하기

여러 개의 정렬된 리스트가 주어졌을 때, 이를 하나의 정렬된 리스트로 병합해야 하는 상황은 실무와 코딩 테스트에서 매우 자주 등장합니다. 이 문제는 힙(heap) 자료구조를 활용하면 효율적으로 해결할 수 있습니다.

예를 들어 [1, 4, 5], [1, 3, 4], [2, 6] 세 개의 정렬된 리스트가 있다면, 병합 결과는 다음과 같습니다.

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

알고리즘 접근 방식

핵심 아이디어는 각 리스트의 첫 번째 원소만 힙에 넣어두고, 가장 작은 값을 꺼낼 때마다 해당 리스트의 다음 원소를 힙에 보충하는 것입니다. 구체적인 단계는 다음과 같습니다.

• n := 리스트 배열 lists의 크기를 저장합니다.
• heap := 새로운 빈 힙(최소 힙)을 생성합니다.
• 각 인덱스 i와 lists[i](row)에 대해:
  • row가 비어있지 않다면, (row[0], i, 0) 튜플을 힙에 삽입합니다.
• res := 결과를 담을 새로운 빈 리스트를 생성합니다.
• 힙이 빌 때까지 반복합니다:
  • num, row, col := 힙의 최상단(최솟값) 요소를 꺼냅니다.
  • res의 끝에 num을 추가합니다.
  • col이 lists[row]의 길이보다 작으면, 즉 아직 남은 원소가 있다면 (lists[row][col + 1], row, col + 1)을 힙에 삽입합니다.
• 병합된 리스트 res를 반환합니다.

시간 복잡도

전체 원소의 개수를 N, 리스트의 개수를 K라고 할 때, 이 알고리즘의 시간 복잡도는 O(N log K)입니다. 모든 원소가 한 번씩 힙에 삽입·삭제되며, 힙의 크기가 최대 K로 유지되기 때문입니다. 이는 모든 원소를 한데 모아 정렬하는 O(N log N) 방식보다 효율적입니다.

예제 코드

아래 구현을 통해 더 잘 이해해 보겠습니다.

import heapq

class Solution:
    def solve(self, lists):
        n = len(lists)
        heap = []
        # 각 리스트의 첫 번째 원소를 힙에 삽입
        for i, row in enumerate(lists):
            if row:
                heapq.heappush(heap, (row[0], i, 0))

        res = []
        # 힙이 빌 때까지 최솟값을 꺼내며 다음 원소를 보충
        while heap:
            num, row, col = heapq.heappop(heap)
            res.append(num)
            if col < len(lists[row]) - 1:
                heapq.heappush(heap, (lists[row][col + 1], row, col + 1))

        return res

ob = Solution()
lists = [[], [], [11, 13], [], [4, 4, 14], [4], [11], [1, 8]]
print(ob.solve(lists))

입력

[[], [], [11, 13], [], [4, 4, 14], [4], [11], [1, 8]]

출력

[1, 4, 4, 4, 8, 11, 11, 13, 14]

마무리

이처럼 Python의 heapq 모듈을 사용하면 정렬된 K개의 리스트를 간결하고 효율적으로 병합할 수 있습니다. 빈 리스트가 포함된 경우에도 초기 삽입 시 조건 검사만으로 안전하게 처리되며, 튜플 (값, 행 인덱스, 열 인덱스)을 함께 저장함으로써 어떤 리스트의 몇 번째 원소인지 추적할 수 있다는 점이 핵심입니다.