[Python] 백준 19638 센티와 마법의 뿅망치 (Heapq)

📌 문제

센티는 마법 도구들을 지니고 여행을 떠나는 것이 취미인 악당이다.

거인의 나라에 도착한 센티는 자신보다 키가 크거나 같은 거인들이 있다는 사실이 마음에 들지 않았다.

센티가 꺼내 들은 마법 도구는 바로 마법의 뿅망치로, 이 뿅망치에 맞은 사람의 키가 ⌊ 뿅망치에 맞은 사람의 키 / 2 ⌋로 변하는 마법 도구이다. 단, 키가 1인 경우 더 줄어들 수가 없어 뿅망치의 영향을 받지 않는다.

하지만 마법의 뿅망치는 횟수 제한이 있다. 그래서 센티는 마법의 뿅망치를 효율적으로 사용하기 위한 전략을 수립했다. 바로 매번 가장 키가 큰 거인 가운데 하나를 때리는 것이다.

과연 센티가 수립한 전략에 맞게 마법의 뿅망치를 이용한다면 거인의 나라의 모든 거인이 센티보다 키가 작도록 할 수 있을까?

입력

첫 번째 줄에는 센티를 제외한 거인의 나라의 인구수 N (1 ≤ N ≤ 105)과 센티의 키를 나타내는 정수 Hcenti (1 ≤ Hcenti ≤ 2 × 109), 마법의 뿅망치의 횟수 제한 T (1 ≤ T ≤ 105)가 빈칸을 사이에 두고 주어진다.

두 번째 줄부터 N개의 줄에 각 거인의 키를 나타내는 정수 H (1 ≤ H ≤ 2 × 109)가 주어진다.

출력

마법의 뿅망치를 센티의 전략대로 이용하여 거인의 나라의 모든 거인이 센티보다 키가 작도록 할 수 있는 경우, 첫 번째 줄에 YES를 출력하고, 두 번째 줄에 마법의 뿅망치를 최소로 사용한 횟수를 출력한다.

마법의 뿅망치를 센티의 전략대로 남은 횟수 전부 이용하고도 거인의 나라에서 센티보다 키가 크거나 같은 거인이 있는 경우, 첫 번째 줄에 NO를 출력하고, 두 번째 줄에 마법의 뿅망치 사용 이후 거인의 나라에서 키가 가장 큰 거인의 키를 출력한다.

예제 입력 1

1 10 1
20

예제 출력 1

NO
10

예제 입력 2

2 10 3
16
32

예제 출력 2

YES
3

예제 입력 3

2 10 3
127
8

예제 출력 3

NO
15

예제 입력 4

1 1 100000
1

예제 출력 4

NO
1

📌 풀이

import sys, heapq
input = sys.stdin.readline

n, h, t = map(int, input().split())
giants = [-int(input()) for _ in range(n)]
heapq.heapify(giants)
cnt = 0

for i in range(t):
    if -giants[0] == 1 or -giants[0] < h:
        break
    else:
        heapq.heapreplace(giants, -(-giants[0] // 2))
        cnt += 1

if -giants[0] >= h:
    print('NO', -giants[0], sep='\n')
else:
    print('YES', cnt, sep='\n')

처음에 브루트포스로 풀었더니 시간초과가 나서(^_ㅠ) heapq로 푼 코드!
우리의 목적은 모든 거인의 키를 센티보다 작게 만드는 것이므로
키가 가장 큰 거인부터 뿅망치를 때려주다가 센티보다 작은 거인에 도달했을 때 뿅망치질을 멈춰주면 되겠다.

heapq

파이썬에서 제공하는 우선순위 큐 모듈. 오름차순을 따른다.

giants = [1, 2, 3, 4, 5, 6]

• heapq.heapify(giants)
  리스트 giants를 힙으로 변환
• heapq.heappush(giants, 7)
  해당 힙(giants)에 원소 7을 푸시
• heapq.heappop(giants)
  해당 힙(giants)에서 가장 작은 원소를 팝
• heapq.heapreplace(giants, 7)
  해당 힙(giants)에서 가장 작은 원소를 팝한 후 원소 7을 푸시
• heaqp.heappushpop(giants, 7)
  해당 힙(giants)에 원소 7을 푸시한 후 가장 작은 원소를 팝

거인들의 키를 오름차순이 아닌 내림차순으로 정렬해야 하므로 힙에 푸시할 때는 -1을 곱해서 푸시하고, 팝할 때 역시 반환된 원소에 -1을 곱해서 연산을 진행한다.

(ex. 거인의 키가 16, 32, 24라고 해 보자. 힙에 -16, -32, -24로 푸시하면 팝할 때나 giant[0]에 접근할 때 가장 작은 -32가 반환되므로 여기에 -1을 곱해서 32로 만들고 우리가 원하는 연산(센티의 키와 비교, 키 반동강 등)을 진행하면 되는 것!)

1. for문은 뿅망치 횟수만큼 돌린다.
2. 현재 가장 큰 거인의 키가 1이면 뿅망치를 때릴 수 없고, 센티보다 작다면 임무를 완수한 것이므로 더이상 뿅망치를 때리지 않아도 된다. → break
3. 현재 가장 큰 거인의 키가 센티보다 크다면 pop하고 뿅망치로 키를 반동강 내 push한다.