[Programmers] [레벨3] 야근지수

참고 : https://programmers.co.kr/learn/courses/30/lessions/12927 

 

이 문제를 풀고나서 다른 사람들과 풀이방식이 조금 다르기에 이렇게 올려본다.

처음에 이 문제를 보았을 때 최댓값을 계속 뽑아내서 -1을 해주면 되겠구나라고 생각을 해서 바로 최대 힙을 생각했다.

파이썬에서는 최대힙을 구현하는 방법은 다들 아시겠지만 원래 숫자에 -1을 곱하면 그렇게 구현할 수 있다.

 

그래서 다음과 같이 구현하였다.

 

import heapq
from functools import reduce
def solution(n, works):
    if n>=sum(works):
        return 0
    heap = []
    for v in works:
        heapq.heappush(heap, -v)
    while n:
        x = heapq.heappop(heap)
        x +=1
        heapq.heappush(heap,x)
        n-=1
    heap = map(lambda x: x**2, heap)
    answer = reduce(lambda x,y : x + y, heap)
    return answer