프로그래머스 level3 - 야근지수

문제 설명

회사원 Demi는 가끔은 야근을 하는데요, 야근을 하면 야근 피로도가 쌓입니다. 야근 피로도는 야근을 시작한 시점에서 남은 일의 작업량을 제곱하여 더한 값입니다. Demi는 N시간 동안 야근 피로도를 최소화하도록 일할 겁니다.Demi가 1시간 동안 작업량 1만큼을 처리할 수 있다고 할 때, 퇴근까지 남은 N 시간과 각 일에 대한 작업량 works에 대해 야근 피로도를 최소화한 값을 리턴하는 함수 solution을 완성해주세요.

제한 사항

works는 길이 1 이상, 20,000 이하인 배열입니다.
works의 원소는 50000 이하인 자연수입니다.
n은 1,000,000 이하인 자연수입니다.

풀이과정

일단 구해야 하는 야근 피로도는? 남은 일들의 작업량 제곱하여 더한 값
1시간에 작업량 1만큼 할 수 있고, n 시간 남아있다고 할 때, 최소 야근 피로도 구하는 문제
여기서 제곱의 합을 줄인다는 것은 가장 큰 값을 줄인다는 것이다.

하지만 조건을 보면 len(works) < 20,000 / work < 50000, n < 1,000,000
만약 n만큼 계속 정렬하면? 시간 복잡도가 어마어마하게 늘어날 것이다.
그래서 list를 정렬하는 것이 아닌 힙을 사용해보자!

힙(heap) 활용하기

• import heapq
• heapq.heappush(heap, item) : item을 heap에 추가
• heapq.heappop(heap) : heap에서 가장 작은 원소를 pop & 리턴. 비어 있는 경우 IndexError가 호출됨.
• heapq.heapify(x) : 리스트 x를 즉각적으로 heap으로 변환함 (in linear time, O(N))

<예시>

import heapq

heap = []
heapq.heappush(heap, 50)
heapq.heappush(heap, 10)
heapq.heappush(heap, 20)

print(heap)

문제에 적용하기

먼저 우리는 work가 큰 것부터 처리하고 싶지만 최소 힙이기 때문에 각 값을 음수로 전환해준다.
어차피 제곱을 하기 때문에 문제되지 않고 일을 했을 때는 -1이 아닌 +1을 해준다.
여기서 0을 넘어가면 +1을 해서는 안된다 (작업량이 0이면 작업을 할 수 없기 때문)

import heapq
def solution(works, n):
    works = [-work for work in works]
    heapq.heapify(works)
    for _ in range(n):
        work = heapq.heappop(works) 
        if work == 0:
            return 0
        else:
            heapq.heappush(works, work+1)
    total = 0
    for work in works:
        total += (work*work)
    return total

works = [4,3,3]
n = 4
print(solution(works, n))