[파이썬] 야근지수 문제

프로그래머스 야근지수(파이썬)

• 프로그래머스 문제 링크

• 힙큐를 이용한 풀이

  import heapq
  
  def solution(n, works):
      for i in range(len(works)): 
          works[i] *= -1
      
      heapq.heapify(works)
      
      for i in range(n):
          m = heapq.heappop(works)
          if m >= 0:
              break
          m += 1
          heapq.heappush(works, m)
      
      answer = 0
      for i in range(len(works)):
          answer += pow(works[i]*-1,2)
      
      return answer
  

  1. 힙큐 모듈을 사용한다. 이 문제의 솔루션은 주어진 works 리스트에서 가장 큰 값을 찾고, n이 0이상일 동안 1씩 빼나가는 것인데, 매번 가장 큰 값을 찾기 위해 소팅하는 것은 비효율적이다. 이 비용을 줄이는 것이 문제의 관건이다.

  2. 힙큐를 사용하면 최소 힙을 구현하여 가장 작은 값을 쉽게 찾을 수 있도록 한다. 모든 원소를 정렬하지 않으므로 소팅하는 것에 비해 효율적이다.

  3. 그러므로 works에 -1을 곱하여, 가장 큰 값이 가장 작은 값이 되도록 한 다음 힙큐를 사용해서 가장 작은 값(원래는 가장 큰 값)을 찾는다.

  4. 이 때 주의할 점은 heapq.heapify()를 사용하여 그냥 리스트였던 works를 힙큐로 바꿔준 다음 heappop을 사용해야 한다는 것이다. 이것을 빼먹고 heappop을 하기 시작했더니 거의 다 틀렸다.

  5. works를 힙큐로 바꾼 다음, n >= 0 일 동안, 힙큐에서 가장 작은 값을 pop해주는 메서드인 heappop을 사용해서 가장 작은 값을 m에 저장한 다음, +1을 하고 이를 다시 힙큐에 넣는다. +1을 하는 것은 원래로 따지면 가장 큰 값에서 -1을 하는 것인데, 지금은 모두 음수인 상태이니 반대로 +1을 하는 것이다. 즉 원래로 따지면 가장 큰 값을 추출하여 여기서 1을 뺀 다음, 다시 그 값을 리스트에 넣어주는 것이다.

  6. 만약 heappop한 값이 0이라면, 모든 야근 작업량이 끝났음을 의미한다. heappop으로 추출된 가장 작은 값이 0이면, 리스트의 모든 값이 0과 같거나 크다는 의미이다. 즉 원래로 따지자면 가장 큰 작업량마저 0이 되었다는 소리니까 모든 작업량이 0이거나 이보다 작다는 소리이다. 하지만 -값이 되기 전에 걸러졌을 것이므로 결국 다 0이라는 소리이다.

     그러면 이제 남은 야근이 없으므로 루프를 종료한다.

  7. 남은 작업량이 없거나, 아니면 n이 0이 되었을 때에 루프가 종료된다. 이 후 works의 각 원소를 제곱하여 답을 도출하면 된다.