Algorithm/Baekjoon
13975. 파일 합치기 3 (Python)
원만사
2022. 2. 18. 14:37
반응형
13975번: 파일 합치기 3
프로그램은 표준 입력에서 입력 데이터를 받는다. 프로그램의 입력은 T개의 테스트 데이터로 이루어져 있는데, T는 입력의 맨 첫 줄에 주어진다.각 테스트 데이터는 두 개의 행으로 주어지는데,
www.acmicpc.net
풀이
파일들을 하나의 파일로 합칠 때의 비용을 계산하기 위한 연산의 횟수는 K-1번이다. 이 때, 최소비용을 구하기 위해서는 각 연산마다 파일 중 제일 작은 값들을 사용해야 다음 연산에 사용할 값을 작은 값으로 만들 수 있다.
최소 힙을 사용하여 두 개의 파일을 선택할 때 현재 파일 목록 중 가장 작은 값을 꺼내오도록 한다. 두 값을 합친 후 다시 최소 힙에 넣어주는 과정을 K-1번 반복한다.
코드
import heapq
def solution(nums):
ans = 0
while len(nums) != 1:
# 최소 힙에서 두 개의 파일을 꺼낸다. 그 후 두 파일을 합쳐서 다시 최소 힙에 넣어준다.
num1 = heapq.heappop(nums)
num2 = heapq.heappop(nums)
tmp = num1 + num2
ans += tmp
heapq.heappush(nums, tmp)
return ans
if __name__ == '__main__':
T = int(input())
for i in range(T):
K = int(input()) # 소설을 구성하는 장의 수
nums = list(map(int, input().split())) # 1장부터 K장까지 수록한 파일의 크기들
pq = [] # 최소 힙
# 입력으로 받은 파일의 크기들을 최소 힙에 넣어준다.
for num in nums:
heapq.heappush(pq, num)
print(solution(pq))
반응형