야근 지수 (Python)

접근 방법

1. n만큼 그리디하게 works에서 1씩 빼주면 최소 야근 지수를 구할 수 있다.
2. 먼저 들어온 works 배열을 힙 구조로 만들어준 뒤, 가장 큰 works 원소부터 1씩 빼준다.
3. n 만큼 소모하면 0 이하의 값을 제외하고 제곱해서 answer에 추가한다.

코드

from heapq import heapify, heappush, heappop

def solution(n, works):
    repl = [-work for work in works]
    heapify(repl)
    for i in range(n):
        heappush(repl, heappop(repl) + 1)
    answer = 0
    for num in repl:
        if num > 0: continue
        answer += num ** 2
    return answer