[백준] 1715번: 카드 정렬하기 (우선순위큐) - Python

https://www.acmicpc.net/problem/1715

코딩 테스트를 준비하면 필수로 알아야하는 문법을 가져왔습니다.

아이디어는 매우 간단한 문제입니다.

---

누가 봐도 중복으로 계산되는 값이 최대한 작게 만드는 문제고

따라서 작은 것들부터 계속 더해가는 그리디로 풀어야하는 것을 알 수 있습니다.

그런데! 여기서 문제가 생겨요.

"작은 것부터 더해? 그러면 정렬을 계속 해야하나?"

계속 더해줘야 해서 무조건 반복문이 들어갑니다. 그래서 O(N)은 강제 할당이고

O(N)을 반복문 안에서 한 번 더 써버리면 무조건 시간초과가 납니다.

문제는 List의 min(), sort() 모두 시간이 O(N)이라는 겁니다.

import sys

n = int(sys.stdin.readline())

cards = []
for _ in range(n):
    number = int(sys.stdin.readline())
    cards.append(number)

answer = 0
while True:
    cards.sort(reverse=True)
    if len(cards) == 1:
        break
    mini_1st = cards.pop()
    mini_2nd = cards.pop()
    answer += mini_1st + mini_2nd
    cards.append(mini_1st + mini_2nd)

print(answer)

그래서 이런식으로 코드를 짜버리면 시간초과가 나와요.

---

그러면 여기서 작은 것들만 정렬 없이 계속 뽑아낼 수 있는 것은 무엇일까요?

바로 바로

"우선 순위 큐" 입니다.

queue 모듈에 PriorityQueue 라는 친구가 있지만

코딩 테스트에서는 더 빠른 동작을 하는 heapq를 사용합니다.

메소드에는

heapq.heappush(heap, item) # O(log n)

item을 heap 리스트에 넣는다.

heapq.heappop(heap) # O(log n)

heap에서 가장 작은 원소를 삭제하고 값을 반환한다.

heapq.heapify(list) # O(n)

list를 heap으로 변환 시킨다.

---

이렇게 알면 됩니다.

그래서 우선순위큐로 구현하면 쉽게 풀려요. 골드4 치고 너무 쉬웠습니당.

import sys
import heapq

n = int(sys.stdin.readline())
cards = []
for _ in range(n):
    number = int(sys.stdin.readline())
    heapq.heappush(cards, number)

answer = 0
while True:
    if len(cards) == 1:
        break
    add = heapq.heappop(cards) + heapq.heappop(cards)
    answer += add
    heapq.heappush(cards, add)

print(answer)

우선순위큐의 디테일은 나중에 다뤄보겠습니다.