[프로그래머스 스터디] 야근 지수

문제

프로그래머스 야근지수

풀이

import heapq


def solution(no, works):
    result = 0
    works = [-i for i in works]

    heapq.heapify(works)

    for i in range(no):
        # 모든 작업을 0으로 만들었음에도 no이 남으면 계속해서 더해지는 예외상황이 존재
        # -> works[0] (가장 큰 값)이 0이 아닐때 조건 추가
        if works[0] != 0:
            task = heapq.heappop(works) + 1
            heapq.heappush(works, task)

    for i in works:
        result += i ** 2

    return result

요약

최대힙을 이용해야하는 문제.
파이썬은 최소힙을 default로 제공하고, 최대힙은 제공해주지 않지만, heapify를 할때 "-"를 붙여줌으로써 최대힙처럼 이용할 수 있다.

1. works = [-i for i in works]
   리스트 Comprehension으로 works리스트를 최대값이 root가 되게끔 heapify함.
2. 주어진 시간만큼 for문을 돌며 아래 작업을 반복.

   작업리스트(works)의 가장 큰 값이 0이 아닐때(해야할 작업이 남아있는 상황)

   모든 작업을 0으로 만들었음에도 no이 남으면 계속해서 더해지는 예외상황이 존재해서 처음에 이부분을 놓쳐 조금 애먹었었다.

• 1. heapq.heappop(works)을 통해 가장 큰값에 +1 ("-"를 붙였기에 일을 줄이기 위해선 더해야함)
• 2. heap에 위 작업을 통해 변경된 남은 작업량을 push한다.

3. 남은 작업량을 각각 제곱해서 result에 더해준후 result return

피드백을 받고 개선된 코드

import heapq


def solution(n, works):
    # 최대힙 구성
    works = [-i for i in works]

    heapq.heapify(works)

    for i in range(n):
        if works[0] != 0:
            heapq.heapreplace(works, works[0] + 1)


    return sum([i ** 2 for i in works])

• heappop, heappush연속될 경우에 heapreplace함수를 이용하면 더욱 효율적이다라는 피드백을 반영
• 제곱한 결과값을 더하는 과정 또한 list comprehension을 통해 간결화