[Level 3] 야근 지수

programmers.co.kr/learn/courses/30/lessons/12927?language=python3

코딩테스트 연습 - 야근 지수

---

접근

초기에 접근했던 방법은 n시간만큼 works에서 최댓값을 찾아서 -1씩 해주는 것이었으며 효율성테스트에서 시간초과가 발생했다.

이후 시간이 많이 남았는데 works는 모두 0일 때 계속 검사를 하게되기 때문에 합이 0일 때는 break를 했지만 시간초과가 발생했다.

 

따라서 어떤 부분에서 최적화를 진행해야할지 생각을 하게되었고, 최댓값을 찾는 연산을 줄이고자 했다. 여기서 우선순위큐가 생각이 났으며 work를 우선순위 큐에 넣어 최댓값이 항상 앞에오도록 하였고, 최댓값이 0일 때 brak하였다.

 

구현

먼저 heapq에 작업량이 많이 남은 것이 앞에오도록 하기위해 튜플 형태로 음수로 짝지어서 넣었다. 이렇게되면 내림차순으로 정렬이된다.

    hq = []
    for work in works:
        heapq.heappush(hq, (-work, work))

 

이제 n번만큼 돌면서 hq의 맨 앞이 0일 때(최댓값이 0일 때) 더 이상 할 일이 없기 때문에 break를 한다.

그리고 최댓값을 가져와서 -1을 해주어 다시 hq에 집어넣는다. 그렇다면 자동적으로 자리를 찾아간다.

for i in range(n):
        if hq[0][1] == 0:
            break
        work = heapq.heappop(hq)
        heapq.heappush(hq, (work[0]+1, work[1]-1))

---

전체 코드

import heapq
def solution(n, works):
    answer = 0
    hq = []
    for work in works:
        heapq.heappush(hq, (-work, work))
    
    for i in range(n):
        if hq[0][1] == 0:
            break
        work = heapq.heappop(hq)
        heapq.heappush(hq, (work[0]+1, work[1]-1))
    
    for item in hq:
        answer += item[1]**2
    return answer