[프로그래머스] 야근 지수 / 파이썬

문제 설명

https://programmers.co.kr/learn/courses/30/lessons/12927

코딩테스트 연습 - 야근 지수

풀이

항상 최대, 최소 값을 접근하면 되니, 힙을 사용하면 되겠구나 라고 쉽게 풀 수 있었던 문제다.
최대힙을 만들기 위해, 음수를 넣었고, 음수를 넣었기 때문에 +1을 더하여 풀었다.
그 외, 일이 사라지는 경우, 힙이 비었을 경우를 따로 처리해주었다.
처리과정이 깔끔하지는 않다.

import heapq

def solution(n, works):
    answer = 0
    h = []
    for w in works:
        heapq.heappush(h, -w)
    for i in range(n):
        if len(h) == 0:
            break
        x = heapq.heappop(h)
        x += 1
        if x != 0:
            heapq.heappush(h, x) 
    for w in h:
        answer += w ** 2
    return answer