문제집 추천, 이 문제집에 나온 유형들만 공부해도 코딩테스트는 거뜬: https://www.acmicpc.net/workbook/by/BaaaaaaaaaaarkingDog
15903번: 카드 합체 놀이 문제보기
문제
석환이는 아기다. 아기 석환이는 자연수가 쓰여져있는 카드를 갖고 다양한 놀이를 하며 노는 것을 좋아한다. 오늘 아기 석환이는 무슨 놀이를 하고 있을까? 바로 카드 합체 놀이이다!
아기 석환이는 자연수가 쓰여진 카드를 n장 갖고 있다. 처음에 i번 카드엔 ai가 쓰여있다. 카드 합체 놀이는 이 카드들을 합체하며 노는 놀이이다. 카드 합체는 다음과 같은 과정으로 이루어진다.
- x번 카드와 y번 카드를 골라 그 두 장에 쓰여진 수를 더한 값을 계산한다. (x ≠ y)
- 계산한 값을 x번 카드와 y번 카드 두 장 모두에 덮어 쓴다.
이 카드 합체를 총 m번 하면 놀이가 끝난다. m번의 합체를 모두 끝낸 뒤, n장의 카드에 쓰여있는 수를 모두 더한 값이 이 놀이의 점수가 된다. 이 점수를 가장 작게 만드는 것이 놀이의 목표이다.
아기 석환이는 수학을 좋아하긴 하지만, 아직 아기이기 때문에 점수를 얼마나 작게 만들 수 있는지를 알 수는 없었다(어른 석환이는 당연히 쉽게 알 수 있다). 그래서 문제 해결 능력이 뛰어난 여러분에게 도움을 요청했다. 만들 수 있는 가장 작은 점수를 계산하는 프로그램을 만들어보자.
입력
첫 번째 줄에 카드의 개수를 나타내는 수 n(2 ≤ n ≤ 1,000)과 카드 합체를 몇 번 하는지를 나타내는 수 m(0 ≤ m ≤ 15×n)이 주어진다.
두 번째 줄에 맨 처음 카드의 상태를 나타내는 n개의 자연수 a1, a2, …, an이 공백으로 구분되어 주어진다. (1 ≤ ai ≤ 1,000,000)
출력
첫 번째 줄에 만들 수 있는 가장 작은 점수를 출력한다.
예제 입력 1
3 1
3 2 6
예제 출력 1
16
예제 입력 2
4 2
4 2 3 1
예제 출력 2
19
아주 간단하지만 푸는동안 찝찝했다. sort를 매 반복마다 수행해야하면 큰 수에 대해서 시간초과가 날거 같았다.
그래서 다른 방법이 있는지 찾아 보았다.
그리고 힙큐를 사용하는 방법이 더 효율적이라는 것을 알게되었다.
처음 풀이, 틀린 코드
전체 카드의 합을 가장 작게 하려면, 카드 더미에서 가장 작은 두 수를 골라 두 수의 합을 덮어쓰는 과정을 반복하면 된다.
이 방법은 카드 2개를 고른 뒤, 다시 넣을 때마다 sort() 함수로 정렬을 수행하는데, sort()의 시간 복잡도는 O(nlogn)이다.
import sys
sys.stdin=open('input.txt','r')
if __name__=="__main__":
n,m=map(int, input().split())
cards=list(map(int, input().split()))
for _ in range(m):
cards.sort()
add=cards[0]+cards[1]
cards[0]=add
cards[1]=add
print(sum(cards))
위 방법보다 2배 이상 시간 효율적인 풀이가 있다.
힙 개념 공부하기
이 방법은 우선순위 큐를 활용하고 있었다. 우선순위 큐는 heap 자료구조를 통해 구현할 수 있다.
우선순위 큐를 사용하면, 값이 바뀐 카드 2개를 다시 우선순위 큐 자료구조 안에 삽입하기만 하면 된다. heap에서 삽입 연산의 시간 복잡도는 O(logn)으로, 훨씬 효율적이다.
파이썬에서는 heapq 내장 모듈을 통해 쉽게 힙을 구현할 수 있다. 이때 주의할 점은, heap이 그 자체로 별도의 자료구조는 아니라는 점이다. heap을 생성하는 대신, 빈 리스트를 생성한 뒤 heapq 모듈을 통해 리스트를 heap으로 사용할 수 있다.
heap 내용 정리
- 리스트 x를 heap으로 변환 : heapq.heapify(x)
- heap 생성 : 빈 리스트 생성 ([])
- item을 heap에 삽입 : heapq.heappush(heap, item)
- heap의 마지막 요소 삭제 : heapq.heappop(heap)
import sys
import heapq as hq
sys.stdin=open('input.txt','r')
if __name__=="__main__":
n,m=map(int, input().split())
cards=list(map(int, input().split()))
hq.heapify(cards)
for _ in range(m):
card1=hq.heappop(cards)
card2=hq.heappop(cards) # 최소값 2개 뽑고
hq.heappush(cards, card1+card2) # 더해서 삽입
hq.heappush(cards, card1+card2) # 알아서 정렬됨
print(sum(cards))