자연수가 쓰여진 카드를 n장 갖고 있다. 카드 합체를 총 m번 하면 놀이가 끝난다. m번의 합체를 모두 끝낸 뒤, n장의 카드에 쓰여있는 수를 모두 더한 값이 이 놀이의 점수가 된다. 만들 수 있는 가장 작은 점수를 계산하는 프로그램을 만들어보자.
# 정답 코드
결국 모두 더한 값이 최소가 되어야 되므로 합체하는 수는 가장 작은 두 수여야 한다.
min()을 이용하면 현재 최솟값은 구할 수 있지만 두번째 최솟값까지 구하려면 최솟값을 제외하고 min()을 한번 더 돌 수 있다. 근데 그렇게 각각 찾으면 index도 따로 찾아서 값을 바꿔줘야 한다. 너무 복잡하다. 최소 O(N^4)
카드 위치를 굳이 고정시켜 놓을 필요는 없잖아?
합체할 때마다 오름차순으로 정렬해주면 항상 맨 앞 두 카드가 가장 작은 두 수다. O(m)*O(NlogN)
* 순서가 필요없을 땐 정렬을 통해 아이디어를 찾는 경우가 많다 !!
import sys
input = sys.stdin.readline
n, m = map(int, input().split())
card = list(map(int, input().split()))
for _ in range(m):
card.sort()
card[0], card[1] = card[0] + card[1], card[0] + card[1]
print(sum(card))# 참고 코드
비교적 쉽게 풀었음에도 이 문제를 블로그에 쓴 이유는 많은 사람들이 heapq를 이용하여 이 문제를 풀었길래 heapq에 대해 알아보려고 한다.
heapq를 사용하면 카드를 뽑을 때마다 정렬하지 않고 heappop()을 2번 해주기만 하면 가장 작은 값 2개가 알아서 뽑힌다. 두 값을 다시 heapq에 넣어주면 끝이다. O(m)*O(logN)
import sys
import heapq
N, M = map(int, sys.stdin.readline().split())
cards = list(map(int, sys.stdin.readline().split()))
heapq.heapify(cards)
for _ in range(M):
tmp = heapq.heappop(cards) + heapq.heappop(cards)
heapq.heappush(cards, tmp)
heapq.heappush(cards, tmp)
print(sum(cards))# heapq
python에서 heap 자료구조는 최소 힙(min heap) 자료구조를 제공한다.
min heap에서 가장 작은 값은 언제나 0번 인덱스(루트)에 위치한다.
min heap 내의 모든 원소(k)는 자식 원소(2k+1, 2k+2)보다 크기가 작거나 같다.
import heapq
card = [] # 빈 리스트를 heap으로 사용하기
heapq.heappush(card, 5); # [5]
heapq.heappush(card, 3); # [3, 5]
heapq.heappush(card, 4); # [3, 4, 5]
heapq.heappush(card, 1); # [1, 3, 4, 5] or [1, 3, 5, 4]
card = [1, 3, 5, 4]
heapq.heapify(card) # 이미 있는 리스트를 heap으로 사용하기
heapq function
heapq.heappush(heap, data) : O(logN)
heapq.heappop(heap) : O(logN)
heapq.heapify(list) : O(N)
