[프로그래머스] 야근 지수 (파이썬)

📕 문제

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

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

 

💡 풀이법

1. 항상 그 순간 가장 작업량이 많은 작업을 해야하기 때문에 계속 큰 값이 나오도록 하는 자료구조인 heap을 이용했다.

2. 최대힙은 없기 때문에 최소힙을 변형해서 최대힙을 만들었다.

3. n까지 for문을 돌면서 매 순간마다 작업량이 최대인 작업을 빼서 작업량이 줄이고 다시 heap에 넣는것을 반복한다.

 

import heapq

def solution(n, works):
    heapq.heapify(works)
    answer = 0
    max_heap = []
    for i in works:
        heapq.heappush(max_heap, (-i, i))
    for _ in range(n):
        if max_heap[0][1] == 0:
            break
        i = heapq.heappop(max_heap)[1]-1
        heapq.heappush(max_heap,(-i,i))
    answer = sum([i[1]**2 for i in max_heap])

    return answer