회사원 Demi는 가끔은 야근을 하는데요, 야근을 하면 야근 피로도가 쌓입니다. 야근 피로도는 야근을 시작한 시점에서 남은 일의 작업량을 제곱하여 더한 값입니다. Demi는 N시간 동안 야근 피로도를 최소화하도록 일할 겁니다.Demi가 1시간 동안 작업량 1만큼을 처리할 수 있다고 할 때, 퇴근까지 남은 N 시간과 각 일에 대한 작업량 works에 대해 야근 피로도를 최소화한 값을 리턴하는 함수 solution을 완성해주세요.
works | n | result |
---|---|---|
[4, 3, 3] | 4 | 12 |
[2, 1, 2] | 1 | 6 |
[1,1] | 3 | 0 |
max heap
를 통해 현재의 최대 작업량에 대해 1씩 처리한다.works
최대힙 heap
으로 구성n
동안 반복하며 최대값 - 1
반복# 코드
import heapq
def solution(n, works):
answer = 0
heap = []
for work in works:
# 최대힙으로 구성하기 위해 (-work, work)의 형태로 push
heapq.heappush(heap, (-work, work))
# n시간 동안 최대 작업량에 대해 1씩 처리
for _ in range(n):
work = heapq.heappop(heap)[1]
# 최대 작업량이 0인 경우
# 모든 작업을 끝낸 것이므로 종료
if work == 0:
break
work -= 1
heapq.heappush(heap, (-work, work))
# 남은 작업량의 각각 제곱하여 더한다.
while heap:
work = heapq.heappop(heap)[1]
if work == 0:
break
answer += (work ** 2)
return answer