프로그래머스 야근 지수 파이썬(PYTHON)

문제링크 : programmers.co.kr/learn/courses/30/lessons/12927

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

 

 

import heapq
def solution(n, works):
    hq=[]
    for i in range(len(works)):
        heapq.heappush(hq,works[i]*-1)
    answer = 0
    for i in range(n):
        if hq[0] == 0:
            return 0
        heapq.heappush(hq,heapq.heappop(hq)+1)
    for i in range(len(hq)):
        hq[i]=hq[i]*hq[i]
    answer = sum(hq)
    return answer