프로그래머스 - LV3. 야근 지수

문제

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

프로그래머스

 

코드

def solution(n, works):
    answer = 0
    
    total_work = sum(works)
    if total_work <= n:
        return 0
    
    for _ in range(n):
        works[works.index(max(works))] -= 1
        
    for w in works:
        answer += w**2 
        
    return answer

 

import heapq

def solution(n, works):
    if sum(works) <= n:
        return 0
    
    works = [-w for w in works]
    heapq.heapify(works)
    for _ in range(n):
        heapq.heappush(works, heapq.heappop(works)+1)
        
    answer = 0
    for w in works:
        answer += w**2 
        
    return answer

리뷰

이 문제의 최종적인 목표는 works 배열 내 모든 원소들의 제곱값의 합을 최소로 만드는 것입니다.

제곱값의 합을 최소로 만드는 핵심적인 방법은, 모든 원소들을 하향평준화 하는 것입니다.

전체적으로 최소값에 가깝게 해야한다는 뜻이죠.

 

이 문제를 푸는데 좋은 모듈로는 heapq 모듈이 있습니다.

쉽게 말해 최소값을 배열의 맨 앞에 유지시켜주는 모듈인데, 이를 여기에 사용하면 쉽습니다.

 

알고리즘은 다음과 같습니다.

• sum(works)가 n 보다 작거나 같을 경우 works 내 모든값이 0이 됨. 
  • return 0
• works를 heapq 배열로 만들어주되, 모든 값에 '-'를 붙임. heapq는 최소값이 우선이기 때문에, '-'를 붙이게 되면 최솟값이 아닌 최대값 기준으로 정렬한 것과 같은 효과를 볼 수 있음.
  • 그냥 heapq 사용할 경우 : [1, 7, 8, 9]
  • '-' 붙여서 heapq 사용할 경우 : [-9, -8, -7, -1]
• 반복문을 통해 n만큼 heap에서 최솟값을 뽑아내고,  1을 더해서 다시 넣어줌.
  • [-9, -8, -7, -1] -> [-8, -8, -7, -1]
• 최종적으로 연산된 works에 대해 제곱합을 하여 반환