[알고리즘] 프로그래머스 - 야근 지수

프로그래머스 - 야근 지수

내 풀이

import heapq

def solution(n, works):
    works = list(map(lambda x: -x, works))
    heapq.heapify(works)
    while n:
        if works[0] < 0:
            work = heapq.heappop(works)
            work += 1
            heapq.heappush(works, work)
        n -= 1
    
    answer = 0
    print(works)
    while works:
        answer += heapq.heappop(works) ** 2
    return answer

파이썬에서 내장함수 heap은 min heap이므로 음수를 곱해서 max heap으로 사용했다. 남은 works에서 가장 큰 값들을 줄이면 된다.