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

https://programmers.co.kr/learn/courses/30/lessons/12927?language=python3 

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

# 프로그래머스 야근 지수
import heapq


def solution(n, works):
    answer = 0
    q = []
    s = sum(works)
    for work in works:
        heapq.heappush(q, -work)
    if n >= s:
        answer = 0
    else:
        for i in range(n):
            work = heapq.heappop(q)
            work += 1
            heapq.heappush(q, work)
        for work in q:
            answer += (work**2)
    return answer

max힙을 이용하여 간단히 구현하였다.