[프로그래머스] 야근 지수(Python)

프로그래머스 - 야근 지수

풀이

리스트에서 가장 큰 수를 가장 작도록 만드는 것이 아니라, 모든 수를 비슷하게 만들어야 한다.
그러기 위해선 리스트에서 큰 수 하나씩 작게 만드는 작업을 여러번 해야하는데, 매번 max값을 찾거나 정렬을 해서 가져오긴 시간이 오래걸리기 때문에 우선순위 큐를 사용했다!

python의 heap을 사용해서 구현했다.
heap은 기본으로 오름차순으로 정렬되기 때문에 편법(?)으로 (-n, n)을 넣어서 -n을 기준으로 정렬을 시켜주면 내림차순 정렬이 된다.

내림차순으로 정렬된 heap에서 하나씩 빼서 -1을 해주고 다시 넣어준다 -> n번 반복

그리고 합을 구하면 정답!

...근데... 다른 사람이 작성한 코드를 봤는데 나처럼 이상한 튜플로 사용하지 않고 처음부터 -n을 넣어서 쉽게 처리했다...

,,,, 하나는 알지만 둘은 모르는,,,,,, 사람,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,

역시 갈 길이 멀다^^

코드

내 코드

import heapq

def solution(n, works):
    max_heap = []
    for i in works:
        heapq.heappush(max_heap, (-i, i))
    while n:
        a, b = heapq.heappop(max_heap)
        if b == 0: return 0
        heapq.heappush(max_heap, (a+1, b-1))
        n -= 1
        
    return sum([b ** 2 for a, b in max_heap])

다른 사람 코드

from heapq import heapify, heappush, heappop
def solution(n, works):
    heapify(works := [-i for i in works])
    for i in range(min(n, abs(sum(works)))):
        heappush(works, heappop(works)+1)
    return sum([i*i for i in works])