[프로그래머스/Lv.3] - 야근 지수(Ref)

문제 링크

코드

아이디어

야근 피로도를 최소로 하기 위하여, 일의 작업량의 최댓값을 일의 작업량의 최솟값으로 만들어야 한다고 생각했다. 이를 위해 먼저 내림차순 정렬(최댓값, 최솟값을 쉽게 찾기 위함)을 수행한다. 그 후, n이 0이 아닐 때 까지 일의 작업량의 최댓값 - (최댓값 - 최솟값)을 수행하고, n에서도 그만큼 빼준다. 이를 수행하면, works의 배열이 내림차순이 아니게 되므로, 다시 내림차순으로 정렬하고 위의 과정을 반복한다.

실패

def solution(n, works):
    if sum(works) <= n: return 0

    while n > 0:
        works.sort(reverse=True)    
        _max, _min = works[0], works[-1]
        diff = _max - _min
        
        if diff <= n:
            works[0] -= diff
            n -= diff
        else:
            works[0] -= n
            break
        
        works.sort(reverse=True)
        _max, _min = works[0], works[-1]
        
        for i in range(len(works)):
            if works[i] == 0: continue
            if n == 0: break
            works[i] -= 1
            n -= 1
            
    return sum(list(work ** 2 for work in works))

일단, 테스트케이스 3개는 통과했지만 나머지는 모두 실패가 떴다. 접근자체가 틀린 것 같진 않은데, 세부적인 로직이 틀린 것 같다.

정답

from heapq import heappop, heappush

def solution(n, works):
    if sum(works) <= n: return 0

    q = []    
    for work in works:
        heappush(q, -work)
    
    for i in range(n):
        _max = -heappop(q)
        if _max <= 0: break
        heappush(q, -(_max - 1))

    result = 0
    for work in q:
        result += work ** 2
        
    return result

다른 코드를 참고하여, 일의 작업량의 최댓값을 찾는 로직을 아예 바꿨다.
먼저, heapq를 이용하여 최대 힙을 만든다. heapq를 사용하면, (최대 - 최소)를 다시 빼는 계산을 하지 않아도 되고 그저 단순히 탐색을 통해 일의 작업량에서 1씩 빼면된다.(n != 0일 때 까지)