처음에는 그냥 전부 합쳐서 한 번에 정렬하면 끝이라고 생각했습니다. 그런데 다시 보니 각 리스트가 이미 정렬되어 있어서, 굳이 전체를 매번 다시 볼 필요는 없었습니다. 결국 이 문제는 단순 정렬 문제가 아니라 현재 후보들 중 가장 작은 값을 빠르게 꺼내는 문제였습니다.
문제
정렬된 리스트가 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
동작 흐름을 손으로 따라가 보기
예시 입력이 아래와 같다고 해보겠습니다.
[
[10, 15, 30],
[12, 15, 20],
[17, 20, 32]
]
처음 힙에는 각 리스트의 첫 값만 들어갑니다.
(10, 0, 0)
(12, 1, 0)
(17, 2, 0)
가장 작은 값 10을 꺼냅니다. 그러면 0번 리스트의 다음 값 15를 다시 힙에 넣습니다.
꺼낸 값: 10
새로 넣는 값: (15, 0, 1)
이제 후보는 이렇게 바뀝니다.
(12, 1, 0)
(15, 0, 1)
(17, 2, 0)
다시 가장 작은 값 12를 꺼내고, 1번 리스트의 다음 값 15를 넣습니다.
꺼낸 값: 12
새로 넣는 값: (15, 1, 1)
이 과정을 계속 반복하면 매 순간 "지금 가능한 후보들 중 가장 작은 값"만 처리하게 됩니다.
이미 각 리스트 안에서는 정렬이 보장되어 있으므로, 리스트 앞쪽 후보만 관리해도 전체 결과가 틀어지지 않습니다.
왜 (value, list_index, element_index)를 넣을까?
힙에 숫자만 넣으면 가장 작은 값은 뽑을 수 있습니다.
하지만 그 값이 어느 리스트의 몇 번째 원소였는지를 모르면 다음 후보를 다시 넣을 수 없습니다.
예를 들어 10을 꺼냈다고 해보겠습니다.
- 0번 리스트의 0번 원소인지
- 1번 리스트의 0번 원소인지
- 2번 리스트의 0번 원소인지
이 정보가 없으면 다음 원소를 이어서 추적할 수 없습니다.
그래서 아래 3가지를 같이 넣습니다.
(value, list_index, element_index)
value: 힙에서 비교할 실제 값list_index: 몇 번째 리스트에서 나왔는지element_index: 그 리스트 안에서 몇 번째 원소였는지
이렇게 해야 가장 작은 값을 꺼낸 뒤, 같은 리스트의 다음 값을 정확하게 다시 넣을 수 있습니다.
실행 예시
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]
실수하기 쉬운 부분
빈 리스트를 그대로 힙에 넣으면 안 됩니다
for list_index, lst in enumerate(lists):
if lst:
heap.append((lst[0], list_index, 0))
여기서 if lst:가 빠지면 빈 리스트에서 lst[0]을 꺼내려다가 바로 에러가 납니다.
값을 꺼낸 뒤 다음 인덱스 범위를 꼭 확인해야 합니다
next_index = element_index + 1
if next_index < len(lists[list_index]):
마지막 원소를 꺼낸 뒤에도 무조건 다음 값을 넣으려고 하면 인덱스 에러가 납니다.
중복 값이 있어도 문제없다는 점을 알아두면 좋습니다
예시에는 15가 두 번 나오고 20도 두 번 나옵니다.
힙은 중복 값을 다룰 수 있으므로 같은 숫자가 여러 리스트에 있어도 괜찮습니다.
오히려 이런 예시를 한 번 직접 따라가 보면, 이 풀이가 값 자체가 아니라 현재 후보들의 순서를 관리하는 방식이라는 점이 더 잘 보입니다.
복잡도
전체 원소 수를 KN이라고 하면:
| 방식 | 시간 복잡도 |
|---|---|
| 단순 정렬 | O(KN log KN) |
| 힙 방식 | O(KN log K) |
리스트 개수 K가 많아질수록 힙 방식이 더 유리합니다. 보조 공간은 힙 크기 기준으로 O(K)입니다.
정리
- 쉬운 방법: 전부 합친 뒤 정렬 —
O(KN log KN) - 더 나은 방법: 힙으로 각 리스트의 현재 후보만 관리 —
O(KN log K) - 핵심: 정렬된 상태 활용, 최솟값 반복 추출, 우선순위 큐
짧아 보이는 문제지만, 설명하는 방식에 따라 면접에서 꽤 차이가 나는 유형입니다.
이 글은 코딩 면접 문제를 바탕으로, 풀이 아이디어를 제 방식대로 재구성해 정리한 글입니다.