프로그래머스 레벨 3 야근 지수

문제

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

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

접근

수학적으로 봤을때, 리스트 내의 최댓값의 수를 줄이는 것이 야근 지수를 가장 많이 줄일 수 있는 방법임 [(n+1)^2-n^2 의 값은 n이 커질수록 커지니까.] 즉, 리스트의 최댓값을 가장 작게 하는 방법을 찾으면 그게 정답임.

그러니까 현재의 최댓값이 몇인지 알아낸후 1씩 감소시켜주는 행위를 반복하면 된다.

매번 리스트를 완전 탐색하여 최댓값을 찾은 후 -1을 하게되면, 최악의 경우 연산횟수가 20000000000가 되어 시간초과가 난다.

즉 완전 탐색을 제외하고 최댓값을 알 수 있는 다른 방법을 찾아야함.

maxheap을 사용하면 풀 수 있다. heap의 시간복잡도는 log2n이고
n=1000000, works의 길이는 최대 20000이므로 총 연산 횟수는 최대
1000000log2(20000)가 된다.

코드

import heapq
def solution(n, works):
    answer = 0
    maxheap=[]
    if sum(works)<=n:
        return answer
    for work in works:
        heapq.heappush(maxheap,-work)
    # heapq 모듈은 기본적으로 minheap 형태이기 때문에 maxheap을
    # 구현하려면 원소의 부호를 바꿔주어야 한다
    while n>0:
        heapq.heappush(maxheap,-(-heapq.heappop(maxheap)-1))
        n-=1
    while maxheap:
        answer+=heapq.heappop(maxheap)**2
    return answer