[Code Tree][Python] 큰 숫자만 계속 고르기

Overview

• 체감 난이도: ★☆☆☆☆
• 문제 유형: 우선순위큐, 힙
• 풀이 상태: 답안참고 / 스스로 해결
• 추후: 다시 풀어보기 / 간단 복습 / 완벽 이해 

---

[문제]

https://www.codetree.ai/missions/8/problems/keep-picking-the-big-number?&utm_source=clipboard&utm_medium=text

코드트리 | 코딩테스트 준비를 위한 알고리즘 정석

 

n개의 숫자가 주어졌을 때 그 중 가장 큰 숫자를 골라 1씩 빼는 작업을 m번 반복하려고 합니다. 이를 반복한 이후 남아있는 숫자들 중 최댓값을 구하는 프로그램을 작성해보세요. 단, 가장 큰 숫자가 여러 개라면 이 중 아무거나 하나를 골라 진행하면 됩니다.

 

 

입력 형식

첫 번째 줄에는 n과 m이 공백을 사이에 두고 주어집니다.

두 번째 줄에는 n개의 숫자가 공백을 사이에 두고 주어집니다.

• 1 ≤ n, m ≤ 100,000
• 1 ≤ 주어지는 숫자들 ≤ 10^9

 

출력 형식

남아있는 숫자들 중 최댓값을 출력합니다. 결과가 음수가 될 수도 있음에 유의합니다.

 

 

[코드]

heapq는 기본적으로 최소 힙 구조이다. 

최대 값을 구하고 싶다면, 음수로 저장하고 pop할 때 음수를 원래대로 돌려준다. 

import heapq

n, m = map(int, input().split())
pq = list(map(lambda x: -int(x), input().split()))  # 입력 값을 음수로 변환하여 최대 힙처럼 사용
heapq.heapify(pq)  # 음수 값으로 최소 힙 구성 -> 최대 힙처럼 동작

for i in range(m):
    x = -heapq.heappop(pq)  # 최대값 꺼내기 (음수로 변환했던 값 복원)
    heapq.heappush(pq, -(x - 1))  # 값을 1 줄인 후 음수로 넣어 최대 힙 유지

# 힙에서 가장 큰 값을 출력 (최대 힙이므로 가장 작은 음수 값)
print(-pq[0])  # `pq[0]`이 가장 작은 음수 값, 즉 가장 큰 양수 값의 음수 형태