프로그래머스 - 야근 지수

문제 설명

야근 지수

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

제한 사항

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

입출력 예worksnresult

[4, 3, 3]	4	12
[2, 1, 2]	1	6
[1,1]	3	0

입출력 예 설명

입출력 예 #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입니다.

입출력 예 #3

 

나의 풀이

import heapq


def solution(n, works):
    if sum(works) < n:
        return 0

    heap = []
    for i in works:
        heapq.heappush(heap, -i)

    for _ in range(n):
        val = heapq.heappop(heap)
        val += 1
        heapq.heappush(heap, val)

    answer = sum(list(map(lambda x: x**2, heap)))

    return answer

우선순위 큐를 생각한 이유는 문제에서 제곱을 하기 때문에 모든 수가 균일하게 작을 때 가장 작은 수가 된다. 때문에 가장 큰 값을 확인해서 계속 줄여나가는 방식을 생각했다. 

 

파이썬에서는 우선순위 큐를 힙이라는 자료구조로 지원하고 있고 최소 트리만을 지원하고 있다. 그래서 최대값을 확인하려 할 때는 (-)음수로 해서 접근한다.

 

1) wokrs 리스트를 heap으로 바꿔준다. 

 

2) n번 만큼 반복문을 실행하고 가장 큰(가장 작은) 값을 계속 더해주며 힙 구조의 상태 변화를 유지하면서 값을 빼주고 더해준다.

 

3) 마지막으로 나온 리스트를 제곱하여 더해주면 끝!

 

 

*처음에는 정렬을 통해 접근을 했으나 반복문이 돌 때마다 정렬을 하는 것은 n**2logn 만큼 시간복잡도가 생긴다. 따라서 우선순위 큐로 힙 자료구조를 통해 시간 복잡도를 줄일 수 있었다.

 

힙의 푸시, 팝 연산은 O(logn) 으로서 최종 시간 복잡도는 O(n*logn)으로 예상한다.