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

문제 접근

• 정렬 된 최대 값이 필요하므로 우선순위큐로 표현해야 된다고 생각하였다
• 가장 빠른 삽입 삭제 시간복잡도 N(데이터) logN을 가진 힙(heap)을 이용하여 문제를 접근

풀이 코드

# 19:10 ~ 19:30
# https://programmers.co.kr/learn/courses/30/lessons/12927?language=python3#
import heapq

def solution(n, works):
    answer = 0
    
    h = []
    
    # 우선순위 큐 최대힙으로 저장
    for value in works:
        heapq.heappush(h, -value)
    
    while n > 0:
        # 힙에서 하나 꺼내어 다시 넣음
        work = -heapq.heappop(h)
        if work != 0:
            heapq.heappush(h, -(work-1))
        else:
            # 일의 작업량의 최대 힙이 0이라는 것은 result도 0이라는 것임
            return 0
        n -= 1
    
    for i in h:
        answer += i**2
    
    return answer