백준 1715 카드 정렬하기 (python)

그리디 알고리즘을 활용하는 문제

 

아이디어

 

 

첫 시도 (시간초과)

n = int(input())
sizes = []
for i in range(n):
    sizes.append(int(input()))

results = sorted(sizes)
answer = 0
while len(results) > 1:
    answer += (results.pop(0) + results.pop(0))
    results.append(answer)
    results = sorted(results)

print(results[0])

리스트를 정렬하고, 가장 작은 값 2개를 꺼내와 더한 뒤 그 값을 다시 추가하는 방식으로 구현해보았다.

결과는 시간초과

 

파이썬 sorted()

파이썬의 sorted()는 Timsort 알고리즘을 사용하여 정렬한다.

Timsort는 병합 정렬과 삽입 정렬을 조합한 알고리즘으로 시간복잡도는 O(n log n)이다.

 

나의 코드는 sorted를 (n-1)번 반복하기 때문에 시간 복잡도가 O(n^2 log n)이 되어버렸다.

n이 최대 10만이므로

10만 * 10만 * log (10만) = 5백억 번 연산이 필요하다.

 

일반적인 컴퓨터에서 1초에 실행할 수 있는 최대 연산이 약 1억 번이라고 하면,

2억 번 연산 안에 풀 수 있는 코드를 작성해야 한다.

 

 

해결 방법 : 우선순위 큐 활용

파이썬의 heapq 모듈을 이용하여 최소 힙 형태의 우선순위 큐를 구현한다.

 

heapq의 heappop(heap)은 heap에서 가장 작은 항목을 반환한다.

 

heapq 모듈을 이용하면

heappop()을 사용하여 가장 작은 값을 O(log n) 시간에 꺼낼 수 있고,

heappush()를 사용하여 값을 삽입하면서 힙 구조를 유지하는 데 O(log n) 시간이 소요된다.

 

import heapq

n = int(input())
sizes = []
for i in range(n):
    heapq.heappush(sizes, int(input()))

answer = 0
while len(sizes) > 1:
    a = heapq.heappop(sizes)
    b = heapq.heappop(sizes)
    answer += a + b
    heapq.heappush(sizes, a + b)

print(answer)

 

heqppush 연산을 for 반복문에서 n번, while 반복문에서 n-1번 반복하고,

heappop 연산을 while 반복문에서 2 * (n-1)번 반복하므로

 

heapq를 사용한 코드의 시간복잡도는 O(n log n)이다.

 

n이 10만일 때,

10만 * log 10만 = 50만 이므로

2초 안에 충분히 해결이 가능해진다.