문제 조건

input

• 일에 대한 작업량 works 배열
  • 1 ≤ works.len ≤ 20,000
• 퇴근까지 남은 시간 n
  • 1 ≤ n ≤ 1,000,000

output

• 퇴근 후 남아있는, works 요소의 최소 제곱함

문제 해석

• 완전 탐색으로 접근 BFS
  • 시간 복잡도 O(n^len(works))
  • 최악의 경우, 20,000 * 1,000,000 → 터짐
• max heap
  • 최소 제곱합을 구하기 위해서는, 현재 works의 최대 값을 순차적으로 줄여줘야 함
  • 최대 값을 탐색하는데 최선의 자료구조는 max heap
  • push, pop 모두 O(logn)
  • n번의 순회 * (pop/push)로 연산을 구체화할 수 있음
  • 최종 시간 복잡도 O(nlogn)

코드

1. DFS

   def solution(n, works):    
   		
   		# 야근을 안 해도 되는 경우 0리턴
       if n >= sum(works):
           return 0
       
       pirodo = float('inf')
       
       def bfs(depth, rest_time, current_pirodo):
           nonlocal pirodo
           
           if depth >= len(works):
               pirodo = min(current_pirodo, pirodo)
               return 
           
           for work_time in range(rest_time + 1): 
               next_pirodo = current_pirodo + (works[depth] - work_time) ** 2
               bfs(depth+1, rest_time - work_time, next_pirodo)
               
       bfs(0, n, 0)
       
       return pirodo if pirodo > 0 else 0
   

2. max heap

   1. n번 loop를 돌면서, works 배열의 최대값을 탐색
   2. 탐색한 값에서 1을 빼주고 다서 works 배열에 push
   3. 파이썬의 heapq 사용
      1. 기본적으로 최소 힙만 지원
      2. 그래서, works 배열에 -1을 곱해줘서 최대 힙으로 사용할 수 있도록 변환

   import heapq
   
   def solution(n, works):
       # 모든 작업을 완료할 수 있는 경우
       if sum(works) <= n:
           return 0
       
       # 최대 힙으로 변환 (Python의 heapq는 최소 힙이므로 음수로 변환)
       works = [-work for work in works]
       heapq.heapify(works)
       
       # n번 작업량 감소
       for _ in range(n):
           # 가장 큰 작업량 꺼내기
           max_work = heapq.heappop(works)
           # 작업량 1 감소 후 다시 힙에 넣기
           heapq.heappush(works, max_work + 1)  # 음수이므로 +1이 감소
       
       # 야근 지수 계산 (남은 작업량의 제곱합)
       # 제너레이터 표현식은 필요한 값을 지연 평가할 수 있어서, 추가 메모리 공간을 사용하지 않음
       # 리스트 컴프리헨션보다 성능상 유리
       return sum(work**2 for work in works)