[Python] 백준/BOJ - 15903번: 카드 합체 놀이

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

15903번: 카드 합체 놀이

 

문제

석환이는 카드 합체 놀이를 하며 놀고 있다. 석환이는 자연수가 쓰여진 카드를 n장 갖고 있다. 카드 합체 놀이는 다음과 같은 과정으로 이루어진다.

1. x번 카드와 y번 카드를 골라 그 두 장에 쓰여진 수를 더한 값을 계산한다. (x != y)

2. 계산한 값을 x번 카드와 y번 카드 두 장 모두에 덮어 쓴다.

이 카드 합체를 총 m번 하면 놀이가 끝난다. m번의 합체를 모두 끝낸 뒤, n장의 카드에 쓰여있는 수를 모두 더한 값이 이 놀이의 점수가 된다. 이 점수를 가장 작게 만드는 것이 놀이의 목표이다.

만들 수 있는 가장 작은 점수를 계산하는 프로그램을 만들어보자.

 

입력

첫 번째 줄: 카드의 개수를 나타내는 수 n(2 <= n <=1000)과 카드 합체를 몇 번 하는지를 나타내는 수 m(0 <= m <= 16 x n)

두 번째 줄: 맨 처음 카드의 상태를 나타내는 n개의 자연수 a1, a2, ..., an이 공백으로 구분되어 주어짐. (1 <= ai <= 1,000,000)

 

출력

첫 번째 줄에 만들 수 있는 가장 작은 점수를 출력한다.

 

예제 입력 1

3 1
3 2 6

 

예제 출력 1

16

 

풀이

이 문제를 풀 때는 n개의 카드 중에서 가장 작은 수의 카드 2개를 뽑아 더한 후 각 카드에 다시 덮어써주면 된다.

 

처음에는 아래와 같이 계속 정렬하는 식으로 풀이를 했었다. 배열을 정렬한 뒤 가장 작은 값 2개 (cards[0]과 cards[1])을 더해주는 식으로 풀었었다. 물론 풀리긴 했지만 이런 식으로 푸는 게 아닌 것 같아서 좀 더 찾아보았다.

n, m = map(int, input().split())
cards = list(map(int, input().split()))

for _ in range(m):
    cards.sort()
    cards[0] += cards[1]
    cards[1] = cards[0]

print(sum(cards))

 

찾아보니 최소힙을 이용하면 되는 것이었다. 파이썬에 heapq라는 모듈을 이용하면 되는 것이었다. 먼저 heapq를 import해주고 카드 n개를 입력받아 cards 리스트에 넣어준다. 그 후 heapq.heapify를 이용해 cards 리스트를 힙으로 변환시켜준 후, m만큼 카드 합체를 반복해주면 된다. heapq는 '최소힙' 자료구조이기 때문에 heapq.heappop(cards)하면 가장 작은 수가 pop해서 나온다. 이렇게 2개의 수를 pop해준 후 그 더한 값을 다시 heapq.heappush를 이용해 힙에 넣어주면 된다. 

 

import heapq

n, m = map(int, input().split())
cards = list(map(int, input().split()))
heapq.heapify(cards)

for _ in range(m):
    a = heapq.heappop(cards)
    b = heapq.heappop(cards)
    heapq.heappush(cards, a+b)
    heapq.heappush(cards, a+b)

print(sum(cards))

 

위: heapq이용, 아래: 노가다 sort..

코드는 좀 더 길어졌지만 실행시간은 훨씬 짧아진 것을 확인할 수 있었다! 파이썬의 유용한 라이브러리들을 열심히 사용해야겠다.