이번 문제는 입력 받은 카드를 오름차순으로 정렬하고 가장 앞의 2개의 수의 합을 덮어 씌우는 과정을 m번 반복하는 방식으로 해결하였다.

처음으로 heapq를 써야겠다고 먼저 생각하게 된 문제였다. 그 이유는 우선 시간 제한이 1초이고, 매 반복마다 문자열을 오름차순 정렬해줘야 한다. m의 범위가 15*n까지 이기 때문에 반복문 자체가 O(n)이고, 매 반복마다 sort()함수를 사용한다면 시간 복잡도가 O(n^2)가 되어 버린다. 이는 당연히 시간초과이다.

heappush()를 이용하면 시간 복잡도는 O(log n)이고, heappop()의 경우는 O(1)이기 때문에 시간 복잡도는 O(n log n)이 된다.

처음에는 일반 배열에 카드 숫자를 입력 받고 for문을 통해 heappush() 하도록 작성하였다. 그러던 중 heapify()가 생각나서 heapify()로 수정하여 다시 작성하였고 이 부분에서의 시간 복잡도 또한 O(n log n) -> O(n)이 되어 수행 시간이 줄어들었다.

• heapq를 불러온다.
• n과 m을 입력받는다.
• 카드 숫자를 저장하는 배열 cards를 선언하고 입력받는다.
• cards를 최소힙으로 바꿔준다.
• m번 반복하는 i에 대한 for문을 돌린다.
  -> 임시 변수 x에 heappop(cards)의 값을 저장한다. 이로 인해 cards의 가장 작은 수는 제거된다.
  -> 임시 변수 y에 heappop(cards)의 값을 저장한다. 이로 인해 cards의 가장 작은 수는 제거된다.
  -> cards에 x+y의 값을 추가한다.
  -> cards에 x+y의 값을 추가한다.
• cards의 총합을 출력한다.

Code

import heapq
n, m=map(int, input().split())
cards=list(map(int, input().split()))
heapq.heapify(cards)
for i in range(m):
    x=heapq.heappop(cards)
    y=heapq.heappop(cards)
    heapq.heappush(cards, x+y)
    heapq.heappush(cards, x+y)
print(sum(cards))