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

Programmers 야근 지수

코드

import heapq

def solution(n, works):
    
    if sum(works) <= n:
        return 0
    
    heap = []
    
    for work in works:
        heapq.heappush(heap, -work)
        
    for _ in range(n):
        max_work = -heapq.heappop(heap)
        max_work -= 1
        heapq.heappush(heap, -max_work)
        
    answer = 0
    
    for work in heap:
        answer += work ** 2
    
    return answer

풀이 과정

처음에는 일반 배열을 이용해서 내림차순 정렬을 반복하여 가장 큰 작업량을 줄이는 순서로 코드를 구현하였다.

def solution(n, works):
    answer = 0
    
    for i in range(n):
        works.sort(reverse=True)
        if works[0] == 0:
            break
        else:
            works[0] -= 1
    
    for work in works:
        answer += work ** 2
    
    return answer

그랬더니 효율성 테스트에서 실패하였다.

위 코드는 매번 works를 내림차순 정렬한 뒤, 가장 앞에 있는 값을 줄인다.

아이디어는 맞지만, n이 크면, n번 반복하면서 매번 전체 배열을 정렬하게 된다. 따라서 매번 최댓값 하나를 찾기 위해 전체 배열을 다시 정렬하는 비효율이 생기게 되었다.

그래서 Heap을 사용하였다.

• Heap: 최댓값이나 최솟값을 빠르게 꺼내기 위한 자료구조

파이썬에서는 heapq 모듈을 사용하며, 기본적으로 최소 힙이므로 가장 작은 값을 먼저 꺼낼 수 있다.

문제에서는 최댓값이 필요하기 때문에 최소 힙만 지원하는 파이썬의 heapq를 활용하기 위해서는 값을 음수로 바꿔서 넣어야했다.

문제에서는 매번 필요한 정보가 ‘현재 가장 큰 작업량’ 하나 뿐이었고, 전체 배열이 매번 완벽하게 정렬되어 있을 필요는 없었다.

처음 풀이처럼 매번 sort()를 하면 전체 배열을 정렬하게 되므로 불필요한 작업이 많아졌다.

반면 힙을 사용하면 매번 최댓값만 빠르게 꺼내고, 다시 넣을 수 있었다.