Merge k Sorted Lists

Q: 정렬된 k개 리스트 병합의 시간 복잡도는 무엇인가요?

최소 힙을 사용하면 O(N log k) 이며, 여기서 N 은 전체 값의 개수이고 k 는 리스트의 개수입니다. 각 값은 한 번씩 삽입되고 제거되며, 힙에는 최대 k 개의 항목이 들어갑니다. 출력 외에 추가로 필요한 메모리는 O(k) 입니다.

Q: 모든 값을 한데 모아 정렬하면 어떨까요?

맞으며, O(N log N) 시간이 걸립니다. 리스트가 이미 정렬되어 있다는 점을 무시하므로 힙이 log k 를 지불하는 대신 각 값마다 log N 을 지불하며, 모든 값을 한꺼번에 메모리에 저장해야 합니다. 힙은 스트림으로 들어오는 리스트도 병합할 수 있습니다.

어려움힙 연결 리스트

+10

lists의 행으로 정수 목록 k개가 주어집니다. 각 행은 비내림차순으로 정렬되어 있고, 행마다 길이가 다를 수 있으며, 비어 있는 행은 없습니다.

모든 행의 모든 값을 포함하고 비내림차순으로 정렬된 하나의 목록으로 병합한 다음 반환하세요. 한 행 또는 여러 행에서 여러 번 나타나는 값은 결과에도 그 횟수만큼 나타납니다.

함수

mergeKLists(lists: integer-2d-array) → integer-array

listsinteger-2d-array: 길이가 서로 다를 수도 있는 정렬된 목록들(행마다 하나씩)
반환값integer-array: 모든 행의 모든 값을 하나의 정렬된 목록으로

제약 조건

1 ≤ lists.length ≤ 10⁴
1 ≤ lists[i].length이고, 모든 행을 합쳐 최대 10⁴개의 값을 가집니다.
-10⁴ ≤ lists[i][j] ≤ 10⁴
모든 행은 비내림차순으로 정렬되어 있습니다.

예제

입력: lists = [[2, 6, 9], [1, 4, 10], [3, 5]]
출력: [1, 2, 3, 4, 5, 6, 9, 10]
설명: 행은 2, 1, 3으로 시작하므로 1이 먼저 옵니다. 그런 다음 각 행의 맨 앞 요소는 2, 4, 3이므로 2가 그다음에 오고, 이런 식으로 계속됩니다. 세 번째 행은 5 다음에 끝나므로 6, 9, 10이 남습니다.

입력: lists = [[5], [-2, 5, 7], [0, 5]]
출력: [-2, 0, 5, 5, 5, 7]
설명: 세 개의 5는 서로 다른 세 행에서 나오며 세 개 모두 남습니다. 음수 -2는 0보다 먼저 정렬됩니다.

후속 질문

각 행의 값을 하나씩 포함하는 가장 작은 범위 [a, b]를 찾으세요. 각 행의 첫 번째 값으로 구성한 힙과 가장 큰 첫 번째 값으로 이를 O(N log k)에 해결할 수 있을까요?

풀이

각 행은 정렬되어 있으므로, 사용되지 않은 값 중 가장 작은 값은 항상 어떤 행의 사용되지 않은 첫 번째 값, 즉 그 행의 맨 앞 값입니다. 모든 k개의 맨 앞 값을 살펴보는 데는 N개의 각 값마다 k단계가 걸립니다. 맨 앞 값들을 담은 최소 힙을 사용하면 가장 작은 값을 O(log k)에 꺼낼 수 있으므로, 전체 시간 복잡도가 O(N·k)에서 O(N log k)로 줄어듭니다.

모든 값에 대해 k개의 헤드를 모두 비교합니다

정답이지만 가장 큰 테스트에서는 끝나지 않음

핵심 아이디어

각 행의 맨 앞에 인덱스 pos[r]를 둡니다. 각 라운드에서는 모든 행의 맨 앞 값을 확인하고, 가장 작은 값을 추가한 뒤 해당 행의 인덱스를 앞으로 이동합니다. 이는 병합 정렬의 병합 단계에서 두 리스트를 k개로 확장한 것입니다. 첫 번째 예시에서는 맨 앞 값이 2, 1, 3으로 시작하므로 1이 출력되고 두 번째 행의 맨 앞 값은 4가 됩니다.

N개의 값을 각각 처리할 때마다 k개 행을 모두 확인해야 합니다. 각 행에 값이 하나씩 있는 행이 10^4개라면, 행을 확인하는 횟수는 10^8회입니다.

알고리즘

모든 행에 대해 pos[r] = 0으로 설정하고 값의 개수인 N을 셉니다.
N번 반복합니다. pos[r]가 아직 해당 행의 범위 안에 있는 모든 행을 살펴보고, 맨 앞 값이 가장 작은 행을 기억합니다.
그 맨 앞 값을 결과에 추가하고 해당 행의 pos에 1을 더합니다.
결과를 반환합니다.

def mergeKLists(lists):
    pos = [0] * len(lists)  # pos[r]: index of the first unused value in row r
    total = sum(len(row) for row in lists)
    merged = []
    for _ in range(total):
        best = -1  # the row whose head is the smallest so far
        for r in range(len(lists)):
            if pos[r] < len(lists[r]) and (best == -1 or lists[r][pos[r]] < lists[best][pos[best]]):
                best = r
        merged.append(lists[best][pos[best]])
        pos[best] += 1
    return merged

k개의 헤드를 담는 최소 힙

핵심 아이디어

각 선택 후 변경되는 헤드는 하나뿐이지만, 스캔할 때마다 전체 k를 다시 읽습니다. 최소 힙은 가장 작은 값이 맨 위에 오도록 헤드들을 유지합니다. 각 헤드에 행을 표시해 두면 (value, row)를 꺼낸 뒤 어느 행의 다음 값을 넣어야 하는지 알 수 있습니다. 힙에는 값이 남아 있는 각 행마다 항목이 하나씩 들어 있으므로, 맨 위 항목이 아직 사용되지 않은 값 중 가장 작습니다.

첫 번째 예에서 힙은 2, 1, 3으로 시작합니다. 1을 꺼내고 4를 넣습니다. 2를 꺼내고 6을 넣습니다. 3을 꺼내고 5를 넣습니다. 4를 꺼내고 10을 넣습니다. 5를 꺼냅니다. 세 번째 행은 모두 사용했으므로 아무것도 넣지 않습니다. 6을 꺼내고 9를 넣은 다음, 9와 10을 꺼냅니다.

각 값은 한 번 들어오고 한 번 나가며, 힙에는 최대 k개의 항목만 들어 있습니다. 따라서 O(log k) 연산을 2N번 수행합니다. N = k = 10^4이면 3 × 10^5단계 미만입니다. 힙은 O(k) 메모리를 사용합니다.

알고리즘

모든 행 r에 대해 (lists[r][0], r)을 값 기준으로 정렬된 최소 힙에 넣습니다.
힙이 비어 있지 않은 동안 가장 작은 쌍 (value, r)을 꺼내고 value를 결과에 추가합니다.
행 r에 다음 값이 있으면 해당 값과 r을 함께 넣습니다.
힙이 비면 결과를 반환합니다.

import heapq


def mergeKLists(lists):
    # The heap holds one (value, row) pair per row that still has values: that row's head.
    heap = [(row[0], r) for r, row in enumerate(lists)]
    heapq.heapify(heap)
    nxt = [1] * len(lists)  # nxt[r]: index of row r's next value, not yet in the heap
    merged = []
    while heap:
        value, r = heapq.heappop(heap)  # the smallest head of all rows
        merged.append(value)
        if nxt[r] < len(lists[r]):
            heapq.heappush(heap, (lists[r][nxt[r]], r))  # row r's new head takes its place
            nxt[r] += 1
    return merged

함정과 경계 사례

값만 넣기. pop한 후에는 어느 행에서 다음 값을 가져와야 하는지 알 수 없습니다. 값과 함께 행도 저장하세요.
실수로 최대 힙을 사용하기. C++ priority_queue와 Rust BinaryHeap은 가장 큰 값을 맨 위에 둡니다. greater<> 또는 Reverse를 사용하세요.
처음에 모든 값을 넣기. 결과는 여전히 정렬되지만, 힙이 N개의 항목까지 커지고 작업량이 O(N log N)이 됩니다.
짧은 행의 끝을 지나서 읽기. 첫 번째 예시에서 세 번째 행에는 값이 두 개 있으므로, 다음 값을 넣기 전에 다음 값이 있는지 확인하세요.
중복 값을 버리기. 두 번째 예시의 5 세 개는 서로 다른 값이며, 모두 결과에 포함되어야 합니다.

자주 묻는 질문3

정렬된 k개 리스트 병합의 시간 복잡도는 무엇인가요?

최소 힙을 사용하면 O(N log k)이며, 여기서 N은 전체 값의 개수이고 k는 리스트의 개수입니다. 각 값은 한 번씩 삽입되고 제거되며, 힙에는 최대 k개의 항목이 들어갑니다. 출력 외에 추가로 필요한 메모리는 O(k)입니다.

모든 값을 한데 모아 정렬하면 어떨까요?

맞으며, O(N log N) 시간이 걸립니다. 리스트가 이미 정렬되어 있다는 점을 무시하므로 힙이 log k를 지불하는 대신 각 값마다 log N을 지불하며, 모든 값을 한꺼번에 메모리에 저장해야 합니다. 힙은 스트림으로 들어오는 리스트도 병합할 수 있습니다.

힙 없이 k개의 정렬된 리스트를 병합할 수 있나요?

네, 분할 정복을 이용하면 됩니다. 두 리스트 병합을 사용해 리스트를 쌍으로 병합한 다음, 결과를 다시 쌍으로 병합하는 식으로 진행합니다. 이렇게 하면 모든 값을 한 번씩 처리하는 라운드가 log k번이므로 O(N log k)입니다. 리스트를 하나씩 늘어나는 결과에 차례로 병합하는 방법은 더 느린 O(N·k)입니다. 초기 값들이 병합할 때마다 다시 복사되기 때문입니다.

def mergeKLists(lists):
    # 여기에 코드를 작성하세요

테스트 #1

테스트 #2

테스트 #3

인자

lists = [[2, 6, 9], [1, 4, 10], [3, 5]]

예상 출력

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