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

문제

야근지수 (https://programmers.co.kr/learn/courses/30/lessons/12927)

풀이

1. 초기 접근

• 문제 조건 해석
  • n은 1,000,000 이하인 자연수입니다.
    • -> 뭔가를 해도, O(n log n) 안으로 끝내야 한다. (for 효율성 테스트)
  • 사실 문제자체는 엄청 간단하다.
    • 직관적으로, 누가 봐도 정렬 후 가장 큰 값만 계속 1씩 빼야될 것 같은데,
      문제는, 리스트 내 가장 큰 값을 n 번 루프동안 어떻게 계속해서 찾아낼 것인가다.
  • 리스트로는 택도없다. 왜냐하면 반복문 돌 때마다 정렬할 것인가?
    • 시간복잡도 루프 O(n) 내, 정렬에 O(n log n) 이 나온다. 따라서 불가능.
• 주목 해야할 것에만 주목하자.
  • 리스트 내 '최대값' 만 필요하다는 것에 주목하자. 리스트로 안된다면, 어떤 자료구조를...?
    • 바로 최대 힙이다.
    • 시간복잡도는, 힙 생성에 O(n log n),
      루프 O(n) 내, 최대 값 뽑는게 O(log n) 이므로, 최대 O(n log n) 안에 풀린다!
  • 파이썬에서 heap 쓰는 것도 배워가자.
• 코드

import heapq

def solution(n, works):
    # 파이썬 기본 heap 은 min heap 이다. 
    #따라서, 모든 원소에 -를 곱해 최소힙으로도 최대힙을 만들 수 있게 하자.
    works = [-work for work in works] 
    heapq.heapify(works)

    for i in range(n):
        work = heapq.heappop(works)
        work += 1
        if work > 0: # 빼먹을 수 있는데, 빼먹으면 안된다.
            work = 0
        heapq.heappush(works, work)

    works = [work*work for work in works]
    return sum(works)