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

나의 코드

import heapq

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

 

풀이 방법

매 연산마다 정렬시 시간 초과, 최소 힙을 이용해 구현