[프로그래머스] Lv.3 야근 지수 [Python, 파이썬]

• 문제 : https://school.programmers.co.kr/learn/courses/30/lessons/12927
• 🔑 최대 힙
  • works의 값들을 최대 힙으로 만든다.
  • 최대 힙은 가장 큰 값이 가장 왼쪽에 있도록 한다.
  • 최대 힙으로 만들려면 값들을 음수로 변경해서 힙에 추가하면 된다.
  • works = [2, 4, 3] 일 때 이 값들을 음수로 바꿔서 힙에 넣으면 (heappush)
  • [-4, -3, -2] 이렇게 가장 큰 값이 왼쪽에 있을 수 있다.
  • heappop을 사용해서 가장 왼쪽 값을 빼낸 다음 1시간 작업한다.
    • 4는 3이 돼야 하는데 현재 값은 -4로 되어있기 때문에 +1을 했다.
  • 1시간 작업을 하고 빼낸 작업을 다시 힙에 넣는다.
  • 이 작업을 n번 반복 후 야근 지수를 구하면 끝 ✨

 

from heapq import heappop, heappush

def solution(n, works):
    answer = 0
    q = []          # 최대 힙
    
    for i in works: 
        heappush(q, -i)         # 최대 힙이기 때문에 음수로 변경해서 힙에 넣음
    
    for _ in range(n):
        x = heappop(q)          # 가장 큰 작업량 뽑기
        if x + 1 < 0:           # 가장 큰 작업을 이번 시간에 끝낼 수 없으면
            heappush(q, x+1)    # 작업 시간 1 처리
        else:                   # 끝낼 수 있으면  
            heappush(q, 0)      # 0 넣기 (끝난 작업)
    
    for i in q:         # 야근 지수 구하기
        answer += i**2
    
    return answer