프로그래머스 level3 야근 지수 (python)

Problem

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

Solution

import heapq

def solution(n, works):
    works = [-work for work in works]
    heapq.heapify(works)
    
    while works and n > 0:
        work = heapq.heappop(works)
        work += 1
        n -= 1
        if work < 0:
            heapq.heappush(works, work)
            
    return sum([work ** 2 for work in works])

각 단계별로 시간이 가장 많이 남은 일을 처리하면 된다. 이를 쉽게 하고자 heap을 사용하였다.