[자료구조] 야근 수당 파이썬 (heap)

from heapq import heapify, heappush, heappop
from collections import deque
def solution(n, works):
    if n > sum(works):
        return 0
    works = [(-i,i) for i in works]
    heapify(works)
    
    for _ in range(n):
        w = heappop(works)[1] - 1
        
        heappush(works, (-w,w))
    ################
    # print works
    # while works:
    #     w = heappop(works)
    #     print(w)
    return sum(i[1]**2 for i in works)
        

 

max value를 하나씩 -1 해줄 것이기 때문에 max heap를 구성해야함

따라서 input으로 (-i, i) 형태로 work를 만들어주고

heapify로 works를 heap으로 구성한다.

 

n번의 시행을 해야하기 때문에 n만큼 for loop

 

할때마다 pop으로 -w가 아니라 w를 꺼내야 하기때문에 [1]을 꺼내서 -1

 

다시 heap에 -w, w 형태로 넣어준다.

----

works의 element의 합이 n보다 작을 경우 음수가 나올 수 있기 때문에 조건을 맨앞에서 처리해준다.

 

++

배상문제와 같은 문제 link