[Programmers] 야근 지수 - Python

 

난이도: Lv. 3 

문제 설명

---

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

 

제한 사항

---

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

 

문제 예제

---

works	n	result
[4, 3, 3]	4	12
[2, 1, 2]	1	6
[1, 1]	3	0

 

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

 

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

 

입출력 예 #3

남은 작업량이 없으므로 피로도는 0입니다.

 

✏️ 내 풀이

---

import heapq

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

    return sum([work ** 2 for work in works])

 

이 문제를 이해하는데 이해하는데 오래 걸린 것 같다.

 

결국, n번 works에서 제일 큰 값을 -1하고 n번 반복이 끝나면 works에 있는 값들을 제곱하여 더해주면 된다.

ex) Input: n=4, works = [4, 3, 3]

• n = 1 --> works = [3, 3, 3]
• n = 2 --> works = [3, 3, 2]
• n = 3 --> works = [3, 2, 2]
• n = 4 --> works = [2, 2, 2]
• result = 2^2 + 2^2 + 2^2 = 12

heapq를 사용하여 이를 쉽게 풀 수 있다.

 

우선, 매 반복마다 가장 큰 값을 꺼내오기 위해 works의 모든 요소에 음수 부호(-)를 붙여준다.

이를 통해 Python의 heapq를 최대 힙처럼 사용할 수 있다.

 

이후 매 반복에서 heappop() 함수를 사용해 가장 큰 값을 꺼내고, 꺼낸 값에 1을 더한 후 heappush()를 통해 다시 works에 삽입한다.

여기서 1을 더하는 이유는, 모든 요소에 음수 부호를 붙였기 때문이다.

 

반복이 끝나면, works에 남아 있는 요소들을 각각 제곱한 후 합산하여 반환하면 결과를 얻을 수 있다.

 

이 문제는 heapq 또는 우선순위 큐에 대한 지식이 없으면 풀기 어려울 수 있지만, 이러한 자료구조를 이해하고 있다면 비교적 쉽게 해결할 수 있는 문제인 것 같다.

---

[Programmers] 야근 지수