처음에 문제 이해가 되지 않아 헤맸다.
풀이를 하자면 N의 수 만큼, works의 가장 큰 원소들을 차례로 -1씩 해주는 과정을 반복하는 것이다.
예를 들어, N이 4일 때는 works의 원소를 총 4번씩 -1 할 수 있는데, 이때 -1을 연산하는 과정은, works의 원소 중 가장 큰 값에 대해서만 적용한다.
문제가 이해되고 나면 접근이 한결 쉬워진다.
works를 오름차순 정렬하면 works[-1]에 대해 -1을 N만큼 반복하면 된다.
import heapq
def solution(N, works):
while N>0:
works.sort()
if works[-1] == 0:
break
works[-1] -= 1
N -= 1 # 굳이 while을 안쓰고 for를 써도 되는 부분
after = list(map(lambda x: x**2, works))
return sum(after)
import heapq
def solution(N, works):
for i in range(N):
works.sort()
works[-1] -= 1
return sum([i**2 for i in works])