문제

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

풀이

시도 1

12938과 비슷한 문제지만 배열의 원소의 합이 최대값이 아닌 제곱의 최소값을 찾아야 하는 문제. 가장 많은 작업을 먼저 처리해야하는 것이 유리하다. heap을 이용한 greedy 알고리즘

  1. works를 힙을 구성하면서 모두 음수로 치환한다. (최대값을 pop 하기 위해, 제곱시 어차피 양수로 변함) - heapq 사용
  2. 작업량을 1 감소시키고 다시 힙에 push
  3. 남은 야근시간 만큼 반복
  4. heap 에 남아있는 원소중 음수는 작업이 필요한 작업량으로 원소에 제곱하여 리턴한다.
solution.py 
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

def solution(n, works):
    import heapq
    # 1 음수로 치환하여 heap을 만든다
    works = [-work for work in works]
    heapq.heapify(works)

    # 2~3 작업량을 1 감소시키고 다시 힙에 푸쉬 (n만큼)
    for i in range(n):
        work = heapq.heappop(works)
        heapq.heappush(works, work + 1)

    # 4 남은것중 음수인 원소는 앞으로 해야할 야근임
    answer = 0
    for i in works:
        if i < 0:
            answer += i*i

    return answer

다른사람풀이

같은 구조지만 파이써닉한 코드

solution.py 
1
2
3
4
5
6

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])