카드 정렬하기

풀이

 이 문제의 핵심 아이디어는 항상 가장 작은 크기의 두 카드 묶음을 합쳤을 때 최적의 해를 보장한다는 것이다.

처음에는 아래와 같이 sort() 함수를 사용해서 정렬시킨 다음 인덱스를 통해 가장 작은 값들을 구했었다. 하지만 시간 초과 판정이 났다. 참고로 아래 코드의 시간 복잡도는 O(n^2logn)이다.

첫번째 풀이(시간 초과)

import sys

N = int(input())

card = []
for i in range(N):
    card.append(int(sys.stdin.readline()))

result = 0
while len(card) != 1:
    card.sort()
    first = card[0]
    second = card[1]
    result += first + second
    card.remove(first)
    card.remove(second)
    card.append(first + second)

print(result)

따라서 우선순위 큐가 떠올랐다. 우선순위 큐는 원소를 넣었다 빼는 것만으로도 정렬된 결과를 얻을 수 있다. 넣었다 빼는 동작은 시간 복잡도가 각각 O(logn)이다. 따라서 아래 코드의 시간 복잡도는 O(nlogn)이다. 우선순위 큐는 힙 자료구조를 이용해서 구현할 수 있다. 파이썬에서는 heapq 라이브러리를 지원하고 있다.

두번째 풀이(정답)

import heapq

N = int(input())
heap = []
for _ in range(N):
    heapq.heappush(heap, int(input()))
result = 0
while len(heap) != 1:
    a = heapq.heappop(heap)
    b = heapq.heappop(heap)
    result += a + b
    heapq.heappush(heap, a + b)
print(result)

시간 복잡도 : O(nlogn)