문제

프로그래머스 야근 지수

문제풀이

n이 주어지고, 배열에서 총 n만큼 뺐을 때 각 원소의 제곱의 합이 최소가 되게 하는 문제이다.

처음에는 프로그래머스 최고의 집합 문제처럼 수학적인 개념을 사용하는 것 같았는데 고민해봐도 마땅한게 떠오르지 않아 컴퓨터를 학대하는 방식으로 방법을 바꿨다.

문제를 푸는데 가장 중요한 점은 배열에서 가장 큰 원소의 값을 줄이면 제곱의 합이 최소가 된다는 것이다.

위의 사항만 숙지하면 문제는 다음과 같은 방식으로 풀린다.

1. n을 정한다. 만약 n이 배열의 총합보다 크면 n을 배열의 총합으로 변경한다.
2. n이 0이 될 때까지 다음의 과정을 반복한다.
   2-1. 배열의 가장 큰 원소를 선택한다.
   2-2. 해당 원소와 n에서 각각 1을 뺀다.

2-1 과정을 수행하기 위해 배열을 최대힙으로 구현하면 정렬이나 max()를 사용해 찾는 것보다 빠르게 $O(logn)$만에 최대값을 찾을 수 있다. 최대힙을 사용하지 않으면 효율성 테스트를 실패한다.

import heapq

def solution(n, works):
    answer = 0

    heap = [(-x, x) for x in works]
    heapq.heapify(heap)
    # 1. init settings
    N = len(works)
    total = sum(works)
    if total - n < 0:
        n = total
    # 2. find answer
    while n > 0:
        _, val = heapq.heappop(heap)
        val -= 1
        heapq.heappush(heap, (-val, val))
        n -=1
    
    res = [x[0]*x[0] for x in heap]
    answer = sum(res)

    return answer

level3 짜리 문제들 중에서는 최대힙만 구현하면 되는 다소(?) 쉬운 문제였던 것 같다.