[자료구조] 배상 비용 최소화 파이썬 (heap)

def solution(no, works):
    
    # max_idx = works.index(max(works))
    # works[works.index(max(works))] -= 1
    while no > 0: # reach 0 exit
        max_idx = works.index(max(works))
        works[max_idx] -= 1
        no -= 1
    
    return sum([x**2 if x > 0 else 0 for x in works])
        
        

return할때 0보다 큰 조건을 안넣어줘서 오답이 나온것 같아서 넣어줬더니 통과했다.

 

일단 여기까지는 했는데

효율성에서 통과를 못했다

---

시간을 줄일 수 있는 부분이 while문 안에서 해결되어야 될거 같은데

음.

 

sorted로 사용하니까 통과하긴 하는데 오래걸리긴 한다.

def solution(n, works):
    while n > 0:
        works = sorted(works)
        works[-1] -= 1
        n -= 1
    return sum([x**2 if x > 0 else 0 for x in works])

---

찾아보니 heapq를 사용해서 우선순위 큐를 사용해야 효율적으로 풀린다.

heap 자체가 최대, 최솟값을 빠르게 찾아내는 자료구조이기 때문에.

 

import heapq
def solution(no, works):
    if no > sum(works):
        return 0
    else:
        works = [(-i,i) for i in works] # max heap
        heapq.heapify(works)
        # print(works)
        
        for _ in range(no):
            work = heapq.heappop(works)[1] -1 #[rank,value] = value
            
            heapq.heappush(works,(-work, work)) # re-push subtracted value to original heap
            
        return sum(i[1]**2 for i in works)

 

기존 sorting보다는 10배 빠르게 실행이 된다.

 

출처: yeahajeong.tistory.com/215

기존 문제: programmers.co.kr/learn/courses/30/lessons/12927