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

---

1. 전체 코드

from math import pow
import heapq

def solution(n, works):
    answer = 0
    h = []
    for i in range(len(works)):
        heapq.heappush(h, -works[i])
    for i in range(n):
        if h:
            val = -heapq.heappop(h)
            heapq.heappush(h, -(val - 1))
    for i in h:
        if i < 0:
            answer += pow(i, 2)
    return int(answer)

---

2. 후기

heapq를 사용하면 풀린다.