[프로그래머스] 야근 지수

수를 골고루 깎아야 최솟값이 나올 것을 유추할 수 있으므로 heapq 모듈을 이용하여 풀었다.

파이썬 기준으로 heapq을 이용하면 자동적으로 최소 힙이 된다. 그러므로 최대 힙으로 바꾸기 위해선 리스트를 이용하여 반대 부호의 값과 함께 넣어 주면 된다.

또한 원래 값이 음수가 되는 경우 제곱하게 되면 양수가 되어 값에 오류가 생기므로 0으로 리셋하여 집어넣는 과정을 거쳤다.

import heapq

def solution(n, works):
    answer = 0
    heap = []
    
    for i in works:
        heapq.heappush(heap, [-i,i])
    
    for i in range(n):
        a,b = heapq.heappop(heap)
        a += 1
        b -= 1
        
        if b < 0:
            b = 0
        
        heapq.heappush(heap, [a,b])
        
    while heap:
        a,b = heapq.heappop(heap)
        answer += b ** 2
        
    return answer