⬛ [Programmers] [Python] [Level3] 야근 지수

문제

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

프로그래머스

풀이

# 프로그래머스 2단계 - 야근 지수
import heapq

def solution(n, works):
    answer = 0
    heap = []
    
    # 시간안에 작업을 끝낼 수 있는 경우
    if sum(works) <= n: 
        return answer

    for work in works:
        # 최대 힙
        heapq.heappush(heap, (-work, work)) 
    
    while n > 0:
        work = heapq.heappop(heap)[1] - 1
        heapq.heappush(heap, (-work, work))
        n -= 1
    for h in heap:
        if h[1] < 0:
            continue
        answer += h[1] ** 2
    return answer

print(solution(1,[2,1,2]))

힙 자료구조를 이용하면 간단히 풀 수 있는 문제이다. 

하지만 구글링을 하기 전까지 힙 자료구조를 사용할 것이라고는 생각지도 못했다. 

최대 힙을 만들고 최대 값을 하나씩 뽑아서 -1 해준 후 다시 힙에 넣어 준다.