Programmers - 야근 지수[파이썬(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 이하인 자연수입니다.

 

입출력 예

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입니다.

 

---

문제 풀이

heap을 활용한 greedy 알고리즘으로 문제를 풀 수 있다.

 

야근 지수는 남은 작업량의 제곱의 합이기 때문에 매 시간 작업량이 가장 많은 작업을 수행하는 것이 가장 야근 지수를 최소로 하는 방법이다. 시간 복잡도 개선을 위해서는 heap을 사용할 수 있는데 남은 시간이 n일 때 O(2n)의 시간 복잡도로 답을 구할 수 있다.(야근 동안 해야 할 작업량을 구하는 과정 O(n), 야근 지수를 구하는 과정 O(n)) n은  1,000,000 이하이므로 최악의 경우  2,000,000번의 연산을 필요로 하므로 효율성을 무리 없이 통과할 수 있다. 

 

python heapq 모듈은 최소 힙으로 구현되어 있으므로 최대 힙을 만들어 주기 위해 각 작업에 음수를 취해 주어야 한다. 답을 구하는 알고리즘은 아래와 같다.

1. 힙에서 작업량이 가장 많이 남은 작업을 pop 한다.
2. 해당 작업량을 1 감소시키고(작업량에 음수를 취했으므로 1을 더해준다.) 다시 힙에 push 한다.
3. 야근까지 남은 시간인 n번만큼 1, 2번을 반복한다.
4. heap에 있는 원소중 음수인 원소는 야근하며 처리해야 하는 작업량이므로 각 원소에 제곱을 전부 더하여 최소 야근 지수를 구한다.

 

구현 코드

코드 링크 : 야근 지수

import heapq

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