문제

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

풀이

제곱을 더한 값을 리턴해야하므로 큰수를 줄이는 것이 작은 수를 0으로 만드는 것보다 효율적이라고 판단했다.
처음에는 일일히 sort해서 큰 수를 pop하고 -1을 해서 다시 넣는 방식을 했더니 당연하게도 효율성테스트에서 걸렸다.
그래서 works의 요소들을 -1을 곱한 최대 힙을 이용해서 풀이를 완료했다.

from heapq import heappush, heappop, heapify
def solution(n, works):
	# 입출력예제 3번처럼 다 더해서 n이랑 작거나 같다면 
    # 일을 완료할 수 있는 것이므로 0 리턴
    if sum(works)<=n:
        return 0
    # 최대 힙 만들기
    works = list(map(lambda x: -x, works))
    heapify(works)
    # 최대 값을 n번만큼 꺼내서 하나씩 줄여간다음에
    # (코드상에서는 최대합이기 때문에 +1임)
    # 다시 힙에 넣는 방식임.
    for i in range(n):
        heappush(works, heappop(works)+1)
    # 람다식을 이용해서 제곱한 값을 sum한 결과물을 리턴
    return sum(map(lambda x : x**2, works))