[프로그래머스] Lv2 - 야근 지수

문제

프로그래머스 연습문제

• Lv 2. 야근 지수 (Python)
  https://school.programmers.co.kr/learn/courses/30/lessons/12927#

코드 구현

import heapq

def solution(n, works):
    answer = 0

    if n >= sum(works):     # 모든 작업을 퇴근까지 남은 시간 안에 다 끝낼 수 있는 경우
        return 0

    works = [-w for w in works]     # works 배열의 모든 값을 음수로 만들어줘서
    heapq.heapify(works)            # 배열 -> 힙 으로 바꿀때 최대 힙을 사용할 수 있도록 함 (heapify == 배열을 힙으로 바꾸는 메서드)

    for _ in range(n):
        i = heapq.heappop(works)    # 힙에서 하나씩 뽑아서
        i += 1      # 1을 빼주는 작업 (음수 + 1 -> 절댓값으로 보면 1이 줄어듦)
        heapq.heappush(works, i)    # 다시 힙에 넣어서 정렬되도록

    # 정답 출력을 위해 각 원소의 제곱을 더해주기
    for i in range(len(works)):
        answer += ((-1) * works[i]) ** 2

    return answer

풀이

• 효율성 테스트에서 애먹은 문제
• 피로도 += 남은 일의 작업량 ** 2 에서 n시간 동안의 피로도를 최소화 해야함 → 즉, 제곱의 합을 최소화 하는 것
• 매 루프마다의 최대값을 선택해서 1씩 빼주면 된다.
• 그러나 sort 함수를 사용하면 시간복잡도 상 효율성 테스트를 통과할 수 없으므로 힙을 사용했다.

초반 코드

from heapq import heappush, heappop

def solution(n, works):
    answer = 0

    # 피로도 += 남은 일의 작업량 ** 2
    # N시간 동안 피로도 최소화
    ## TODO : 매번의 최댓값을 구한다 -> 해당 최댓값에서 -1씩 해줌

    work_heap = []
    for num in works:
        heappush(work_heap, -num)

    while (n > 0):
        to_sub = (-1) * heappop(work_heap)
        if (to_sub == 0):
            break
        elif (to_sub < n):  # 최대값이 n보다 작은 경우, n을 한번에 줄이기
            n -= to_sub
        elif (to_sub >= n):  # 1씩 빼주는 경우
            to_sub -= 1
            n -= 1
            heappush(work_heap, -to_sub)

    for i in range(len(work_heap)):
        answer += ((-1) * work_heap[i]) ** 2

    return answer

→ 예외가 존재하였다. 아래 테스트케이스의 경우, 580이 나오지 않고 더 큰 값이 나옴

solution(99, [2, 15, 22, 55, 55])     # 580

What I learned

파이썬 자료구조 힙 사용법

• heapq 모듈은 이진 트리(binary tree) 기반의 최소 힙(min heap) 자료구조를 제공
  참고 : https://www.daleseo.com/python-heapq/

▶️ heapify() : 기존 리스트를 힙으로 변경

from heapq import heapify

arr = [4, 1, 7, 3, 8, 5]
heapify(arr)

print(arr)
# [1, 3, 5, 4, 8, 7]

▶️ heappush() : 힙에 원소 추가 (리스트의 append와 동일)

▶️ heappop() : 힙에서 원소 제거 (리스트의 pop과 동일)
→ 리스트와의 차이점 : 추가, 제거가 일어날 때 정렬이 이루어짐

▶️ 최대힙 사용하기 : 단순하게 모든 수를 반대로 (양수 → 음수, 음수 → 양수) 로 만들어주면 음수에서 최소~최대 정렬이 되기 때문에 가장 앞에 오는 원소가 최대값이 된다.

:= 사용법

from heapq import heapify, heappush, heappop

heapify(works := [-i for i in works])

• 한번에 works 배열을 최대힙으로 바꾸는 방법