[백준 1715번] 카드 정렬하기(파이썬)

문제

정렬된 두 묶음의 숫자 카드가 있다고 하자. 각 묶음의 카드의 수를 A, B라 하면 보통 두 묶음을 합쳐서 하나로 만드는 데에는 A+B 번의 비교를 해야 한다. 이를테면, 20장의 숫자 카드 묶음과 30장의 숫자 카드 묶음을 합치려면 50번의 비교가 필요하다.

매우 많은 숫자 카드 묶음이 책상 위에 놓여 있다. 이들을 두 묶음씩 골라 서로 합쳐나간다면, 고르는 순서에 따라서 비교 횟수가 매우 달라진다. 예를 들어 10장, 20장, 40장의 묶음이 있다면 10장과 20장을 합친 뒤, 합친 30장 묶음과 40장을 합친다면 (10 + 20) + (30 + 40) = 100번의 비교가 필요하다. 그러나 10장과 40장을 합친 뒤, 합친 50장 묶음과 20장을 합친다면 (10 + 40) + (50 + 20) = 120 번의 비교가 필요하므로 덜 효율적인 방법이다.

N개의 숫자 카드 묶음의 각각의 크기가 주어질 때, 최소한 몇 번의 비교가 필요한지를 구하는 프로그램을 작성하시오.

 

입력

첫째 줄에 N이 주어진다. (1 ≤ N ≤ 100,000) 이어서 N개의 줄에 걸쳐 숫자 카드 묶음의 각각의 크기가 주어진다. 숫자 카드 묶음의 크기는 1,000보다 작거나 같은 양의 정수이다.

출력

첫째 줄에 최소 비교 횟수를 출력한다.

 

예제 입력 1

3
10
20
40

예제 출력 1

100

 

---

 

풀이

이 문제는 두 묶음의 카드를 합쳐서 하나로 만드는데 A+B 번의 비교를 해야되는데, 이러한 비교가 최소한 몇 번 필요한지 구하는 문제입니다.
최소한 비교를 하기 위해선 적은 숫자부터 차례로 더해서 비교해주면 됩니다.
하지만 일반적인 정렬로 처했을 땐 시간 초과가 발생하네요.
그래서 heapq 모듈을 사용하여 시간 초과가 발생하지 않도록 해줍니다.

heapq.heapify() : 리스트 내부의 원소들을 힙 구조에 맞게 재배치되며 최소값이 0번째 인덱스에 위치 하게 됩니다.
heapq.heappop() : 힙에서 가장 작은 원소를 삭제 후 그 값을 리턴합니다.
heapq.heappush() : 힙에 원소를 추가합니다. 첫번째 인자는 원소를 추가할 대상 리스트, 두번째 인자는 추가할 원소이다.

예제 1
3
10
20
40
(10+20) + (30+40) = 100
10+20, 40 / 10+20= 30
30 + 40 / 30+40= 70
70 / 70+0= 70
30 + 70 = 100

코드를 순서대로 작성한다면
1. 최소 힙을 생성 후 카드 묶음 입력
2. 입력된 카드 묶음을 재배치
3. 힙 안에 원소가 1개가 될 때까지 반복
4. 최소값 원소 2개를 꺼내서 더하기
5. 더한 값을 힙 안에 저장
6. 더한 값을 변수에 저장

※ sys.stdin.readline 을 쓰는 이유는 입력 시간 초과 때문입니다.
   한 두개의 입력일 때에는 input을 사용하고, 문제처럼 입력이 많을 때 sys.stdin.readline을 사용합니다.

---

작성 코드

import sys, heapq

N = int(input())

data = [int(sys.stdin.readline()) for i in range(N)] # 힙 생성, 카드 묶음 입력
heapq.heapify(data) # 힙 재배치

answer = 0

while len(data) > 1 : # 힙 안에 원소가 1개일 때까지 반복
    sum = heapq.heappop(data) + heapq.heappop(data)
    heapq.heappush(data, sum) # 힙 안에 더한값을 추가
    answer += sum # 더한 값은 정답 변수에 저장
    
print(answer)