처음에는 그냥 전부 합쳐서 한 번에 정렬하면 끝이라고 생각했습니다. 그런데 다시 보니 각 리스트가 이미 정렬되어 있어서, 굳이 전체를 매번 다시 볼 필요는 없었습니다. 결국 이 문제는 단순 정렬 문제가 아니라 현재 후보들 중 가장 작은 값을 빠르게 꺼내는 문제였습니다.

문제

정렬된 리스트가 K개 주어집니다. 이를 하나로 합쳐서 정렬된 리스트를 반환하면 됩니다.

[
    [10, 15, 30],
    [12, 15, 20],
    [17, 20, 32]
]

결과:

[10, 12, 15, 15, 17, 20, 20, 30, 32]

처음 떠오르는 풀이

def merge(lists):
    arr = []
    for lst in lists:
        arr.extend(lst)
    return sorted(arr)

구현도 쉽고 결과도 맞습니다. 하지만 각 리스트가 이미 정렬되어 있다는 조건을 전혀 활용하지 못합니다.

시간 복잡도는 O(KN log KN).

진짜 봐야 하는 점

매 순간 중요한 건 딱 하나입니다. 각 리스트의 맨 앞 후보들 중에서 가장 작은 값이 무엇인지 보는 것입니다.

처음 상태에서 비교할 값은 사실 세 개뿐입니다.

- 첫 번째 리스트: 10
- 두 번째 리스트: 12
- 세 번째 리스트: 17

10이 가장 작으니 먼저 꺼내면 됩니다. 그러면 첫 번째 리스트의 다음 값인 15가 새 후보가 됩니다.

- 첫 번째 리스트: 15  ← 방금 꺼낸 자리에 다음 값이 들어옴
- 두 번째 리스트: 12
- 세 번째 리스트: 17

이 과정을 반복하면 전체 정렬 결과를 얻을 수 있습니다. 이미 지나간 값보다 더 작은 값이 뒤에서 갑자기 튀어나올 일은 없기 때문입니다.

즉, 핵심은 현재 후보들 중 최솟값을 빠르게 꺼내는 것입니다. 이럴 때 잘 맞는 자료구조가 힙(heap)입니다.

왜 힙을 쓸까?

힙은 여러 값 중 가장 작은 값을 O(log K)에 꺼낼 수 있는 구조입니다.

힙에 넣을 정보는 이렇게 구성합니다.

(value, list_index, element_index)

예를 들어 (10, 0, 0)이면 0번 리스트의 0번 인덱스 값 10이라는 뜻입니다. 가장 작은 값을 꺼낸 뒤 어느 리스트에서 나왔는지 알 수 있으므로, 그 리스트의 다음 원소를 바로 힙에 넣을 수 있습니다.

코드

아래 코드는 위 흐름을 그대로 옮긴 구현입니다. 힙에는 각 리스트의 현재 원소만 넣고, 가장 작은 값을 꺼낼 때마다 해당 리스트의 다음 원소를 다시 넣습니다.

import heapq

def merge(lists):
    merged = []
    heap = []

    for list_index, lst in enumerate(lists):
        if lst:
            heap.append((lst[0], list_index, 0))
    heapq.heapify(heap)

    while heap:
        value, list_index, element_index = heapq.heappop(heap)
        merged.append(value)

        next_index = element_index + 1
        if next_index < len(lists[list_index]):
            next_value = lists[list_index][next_index]
            heapq.heappush(heap, (next_value, list_index, next_index))

    return merged

실행 예시

print(merge([[10, 15, 30], [12, 15, 20], [17, 20, 32]]))
# [10, 12, 15, 15, 17, 20, 20, 30, 32]

print(merge([]))                        # []
print(merge([[], [], []]))              # []
print(merge([[], [1], [1, 2]]))         # [1, 1, 2]
print(merge([[1], [1, 3, 5], [1, 10, 20, 30, 40]]))
# [1, 1, 1, 3, 5, 10, 20, 30, 40]

복잡도

전체 원소 수를 KN이라고 하면:

방식시간 복잡도
단순 정렬O(KN log KN)
힙 방식O(KN log K)

리스트 개수 K가 많아질수록 힙 방식이 더 유리합니다. 보조 공간은 힙 크기 기준으로 O(K)입니다.

정리

  • 쉬운 방법: 전부 합친 뒤 정렬 — O(KN log KN)
  • 더 나은 방법: 힙으로 각 리스트의 현재 후보만 관리 — O(KN log K)
  • 핵심: 정렬된 상태 활용, 최솟값 반복 추출, 우선순위 큐

짧아 보이는 문제지만, 설명하는 방식에 따라 면접에서 꽤 차이가 나는 유형입니다.

이 글은 코딩 면접 문제를 바탕으로, 풀이 아이디어를 제 방식대로 재구성해 정리한 글입니다.