[프로그래머스 | Python] 야근 지수

프로그래머스_야근 지수

제곱의 합이 최소가 되기 위해서는 최댓값을 줄여야하므로, works의 최댓값을 1씩 줄여나간다

정확성 통과, 효율성 실패 답안

def solution(n, works):
    answer = 0
    if sum(works)<=n:
        return answer
    
    while n:
        idx = works.index(max(works))
        works[idx] -= 1
        n -= 1
    
    for w in works:
        answer += w**2
    
    return answer

문제는 n의 최댓값이 1e6이라는 것
그래서 위의 코드는 효율성에 실패했다
heapq를 떠올리기는 했는데, 어떻게 활용해야할지 모르겠어서 결국 다른 분의 코드를 확인했다

둘 다 통과한 답안

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

기존의 heapq는 최소힙이다
즉 최솟값을 먼저 pop 시키는데, 우리는 최댓값에다가 연산을 해야하기 때문에 -를 붙인 최대힙을 만들어준다
그리고 마지막에 제곱의 합을 구할 때 다시 -를 붙여줄 필요가 없는게, 음수를 제곱하면 결국 양수가 되기 때문에 그대로 제곱하여 더해주면 된다