[Algorithm] Programmers : 야근 지수 by Python

[문제 바로가기] https://programmers.co.kr/learn/courses/30/lessons/12927

📌문제 설명

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

제한 사항

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

입출력 예

입출력 예 설명
입출력 예 #1
n=4 일 때, 남은 일의 작업량이 [4, 3, 3] 이라면 야근 지수를 최소화하기 위해 4시간동안 일을 한 결과는 [2, 2, 2]입니다. 이 때 야근 지수는 22 + 22 + 22 = 12 입니다.

입출력 예 #2
n=1일 때, 남은 일의 작업량이 [2,1,2]라면 야근 지수를 최소화하기 위해 1시간동안 일을 한 결과는 [1,1,2]입니다. 야근지수는 12 + 12 + 22 = 6입니다.

💡 문제 풀이

문제내용이 간단하고 제한사항에서 주어진 인자(works, n)의 범위가 큰 것으로 보아 시간복잡도를 해결하는 것이 중요하다고 파악했다.

'야근 피로도'는 남은 일의 작업량을 제곱하여 더한 값이기 최소화한 값을 구하려면 works 배열의 최대값부터 차례대로 최소화해야 한다.

배열내의 최대값의 위치를 찾기 위해서 max, index함수를 사용하였는데 해당 함수를 사용하면 배열을 모두 탐색하는 시간 복잡도를 가지기 때문에
최악의 경우(n은 1,000,000 / works 길이는 20,000)일 때 엄청난 시간이 걸리게 된다.

따라서 적절한 자료구조의 사용이 필요함을 느꼈고 정렬을 할 필요 없이 최소값(최소힙)을 얻을 수 있는 heapq 모듈을 사용하여 효율성을 최적화했다.

step 1)
먼저, 남은 일의 작업량의 총합이 야근 가능시간(n)보다 작게되면 모든 일을 끝 마칠 수 있기 때문에 0을 바로 return 할 수 있도록 조건문을 만든다.

step 2)
heapq에서 pop을 하면 최소값부터 순서대로 빼낸다.
하지만, 문제를 풀기 위해서 우선적으로 필요한건 works 배열의 최대값이기 때문에 배열내의 값을 음수 값으로 바꾼다.(순서를 역순으로 바꾸기 위한 작업)

step 3)
works를 heap으로 사용할 것이기 때문에 heapify 해준다.

이후 야근 가능시간(n)이 0이 될 때 까지 heap에서 최소값을 빼내어 +1을 해주고 동시에 n값을 1씩 감소시킨다.
= 실제로는 가장 오래걸리는 시간-1을 해야되지만 음수로 바꾸었기 때문에 +1을 해준다.

step 4)
야근 피로도는 남은 일의 작업량을 제곱하여 더한 값이므로 works 배열 내의 값들을 제곱하여 answer에 더해준다.
음수의 제곱은 양수이므로 굳이 양수로 바꾸는 작업을 할 필요가 없다.

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

이미 존재하는 배열을 heap으로 사용하기 위해서는 heapify 함수를 사용해야하는데... 생략하고 heappop, heappush만을 사용해서 틀린 이유를 찾지 못했었다.

원소가 있는 배열을 heap으로 사용할 경우에는 반드시 heapify를 해주자!!!