[programmers] [Level3] 야근지수

https://programmers.co.kr/learn/courses/30/lessons/12927

코딩테스트 연습 - 야근 지수

풀이

• 야근 시간 n시간, 처리해야할 업무 리스트 works를 입력으로 받음.
• 야근 지수 = n시간동안 works의 일을 처리(숫자 감소) 후 남아있는 원소의 제곱 합
• n시간동안 매번 큰 수를 찾아서 1씩 감소시킨 후 제곱합을 구하면 됨.
• 리스트의 내장함수 [].index(element)는 O(n)이기 때문에 시간초과 남.
  • "리스트의 max value 찾음 -> 1 감소"를 반복하기 때문에 heapq를 사용해 O(n \log n) 으로 복잡도 줄임.

def solution(n, works):
    import heapq

    if sum(works) <= n :
        return 0

    h = []
    for work in works:
        # 최솟값을 제일 앞으로 저장하기 때문에 (-1)를 곱한 값을 heapq에 저장
        heapq.heappush(h, (-1) * work)

    for i in range(n):
        num = heapq.heappop(h)
        # 1 감소시킨 뒤 다시 넣음
        num = (num * (-1) - 1) 
        heapq.heappush(h, (-1) * num)

    answer = 0
    for i in h:
        answer += (i) * (i)

    return answer