[프로그래머스] 야근 지수 - 파이썬(python)

프로그래머스

문제

회사원 Demi는 가끔은 야근을 하는데요, 야근을 하면 야근 피로도가 쌓입니다. 야근 피로도는 야근을 시작한 시점에서 남은 일의 작업량을 제곱하여 더한 값입니다. Demi는 N시간 동안 야근 피로도를 최소화하도록 일할 겁니다.Demi가 1시간 동안 작업량 1만큼을 처리할 수 있다고 할 때, 퇴근까지 남은 N 시간과 각 일에 대한 작업량 works에 대해 야근 피로도를 최소화한 값을 리턴하는 함수 solution을 완성해주세요.

제한 사항

• works는 길이 1 이상, 20,000 이하인 배열입니다.
• works의 원소는 50000 이하인 자연수입니다.
• n은 1,000,000 이하인 자연수입니다.

풀이 설계

• 야근 지수를 줄이기 위해서는 큰 수가 있으면 안된다.
  • [4,0]보다는 [2,2]가 좋은 방식이다.
• sort(reverse=True)를 사용해 큰수부터 가져올 수 있도록 정렬한다.
• while문을 돌리며 원소를 하나씩 가져와서 -1하고 다시 정렬한다.
• 그러다가 n이 0이 되면 남은 리스트를 for문으로 원소를 제곱해서 더한다.

1차 시도 (시간 초과)

• 매번 sort로 정렬해서 그런지 시간초과가 발생했다.

def solution(n, works):
    answer = 0
    works.sort(reverse=True)
    while n > 0:
        tm = works.pop(0)
        tm -= 1
        works.append(tm)
        works.sort(reverse=True)
    for work in works:
        answer += work**2
    return answer

2차 시도 (성공)

• 정렬로 인한 시간을 줄이기 위해 힙을 사용했다.

import heapq

def solution(n, works):
    answer = 0
    works_heap = []
    for work in works:
        heapq.heappush(works_heap,-work)
    while n > 0:
        tm = -heapq.heappop(works_heap)
        if tm == 0: break
        tm -= 1
        n -= 1
        heapq.heappush(works_heap,-tm)
    for work in works_heap:
        answer += work**2
    return answer

효율성 테스트 결과

후기

• 참고할만한 다른 풀이가 모두 힙을 사용해서 다른 풀이는 생략했다.
• heapq를 안다면 쉽게 풀 수 있는 문제였다.