[프로그래머스] LV.3 야근 지수 (Python) - heapq

 

 

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

프로그래머스

max heap을 사용하는 문제. 파이썬에서 max heap을 구현할 때는 우선순위에 -를 붙여준다고 생각하면 된다. 

 

 

 

제출한 코드

import heapq
def solution(n, works):
	remain_time = min(sum(works), n) # 몇시간 일하는지
    # max heap 만들기
    heap = [-work for work in works]
    heapq.heapify(heap)
    while True:
        if remain_time <= 0:
            break
        i = heapq.heappop(heap)
        # 마이너스 붙인 상태로 정렬된 상태니까 원소에 1 더해주면 -4 -> -3이 되어 우선순위에서 밀려날 수 있음
        # 이 문제의 경우, 마지막에 제곱하기 때문에 그대로 음수로 둬도 됨
        #heap의 원소를 업데이트하려면 heappop으로 꺼낸 뒤, 업데이트하고 다시 push해야 함
        i += 1
        heapq.heappush(heap, i)
        remain_time -= 1

    return sum(work**2 for work in heap)

 

 

 

🔥주의점

min heap은 맨 앞 원소가 최솟값인 것만 보장한다. heap[0]으로 가장 작은 원소에 접근하는 것은 가능하지만, heap[1], heap[2], ... 순서대로 점점 커지는지는 보장할 수 없다. 

 

 

✨ 이 문제를 풀면서 다시 외운 것들

## sort의 시간복잡도는 nlogn
## heapify의 시간복잡도는 nlogn (heappop, heappush는 logn)