[Heap] 야근 지수 - Python / JavaScript

문제 링크 - https://school.programmers.co.kr/learn/courses/30/lessons/12927

프로그래머스

 

[문제 간단 요약]

• 남은 시간 n, 각 작업의 작업량 배열 works가 주어짐
• 1시간 동안 한 작업량을 1 줄일 수 있음
• 모든 시간이 끝난 뒤 남은 작업량들의 제곱합이 최소가 되도록 작업을 분배
• 결과: 최소화된 제곱합

 

[문제 해결 방안]

• 제곱합을 최소화하려면 가장 큰 작업부터 줄여야 함
• 따라서 최대 힙 구조를 사용해야 함
• Python: heapq에 음수를 넣어 최대 힙처럼 사용
• JavaScript: 힙 내장 기능이 없으므로 MaxHeap 클래스를 직접 구현

[문제 해결 코드 - python]

import heapq
def solution(n, works):
    answer = 0
    hq = []
    for work in works:
        hq.append(-work)
    heapq.heapify(hq)
    while True:
        if n == 0 or len(hq) == 0:break
        n-=1
        pop_num = heapq.heappop(hq)+1
        if pop_num == 0:continue
        heapq.heappush(hq, pop_num)
    for h in hq:
        answer += h**2
    return answer

[코드 개선점]

• 불필요한 조건 단순화
  Python 코드의 if pop_num == 0: continue는 if pop_num != 0: push 형태로 가독성 향상 가능
• 시간 복잡도
  삽입/삭제는 모두 O(log n), 최대 연산 횟수는 n → 전체 O(n log m), (m = 작업 개수)
• 공간 복잡도
  힙에 최대 m개의 원소 저장하므로 O(m)
• 예외 처리
  작업량 총합 ≤ n이면 자동으로 [0,0,…] → 제곱합 0 반환

[개선된 코드 - python]

import heapq

def solution(n, works):
    hq = [-w for w in works]
    heapq.heapify(hq)

    while n > 0 and hq:
        largest = -heapq.heappop(hq)
        if largest > 1:
            heapq.heappush(hq, -(largest - 1))
        elif largest == 1:
            pass
        n -= 1
	return sum(val * val for val in hq)