코딩테스트 문제 (28) - 정렬 리스트 병합하기

문제

여러 개의 정렬된 리스트를 1개의 정렬된 리스트로 병합하는 문제입니다.

a = [[1,4,5],[1,3,4],[2,6]]
=> [1,1,2,3,4,4,5,6]

 

풀이

이 문제는 우선순위 큐를 사용해 풀 수 있는 문제로 파이썬에서 heapq 모듈을 사용하면 됩니다. 파이썬의 heapq 모듈은 최소힙이 구현된 것으로 리스트 원소를 heappush 메소드로 넣고 heappop 메소드로 리스트를 추출할때 이미 정렬되어 있으므로 맨 앞의 원소를 뽑아내고 다시 heap 구조에 집어 넣습니다.

from heapq import heappush, heappop

def solution(ListofList):
    result = []
    heap = []
    for i, l in enumerate(ListofList):
        heappush(heap, (l))

    print(heap)

    while heap:
        l = heappop(heap)
        result.append(l.pop(0))
        if len(l) > 0:
            heappush(heap, (l))
    
    return result

Result