Programmers / 야근 지수 / Python

*문제 출처는 프로그래머스에 있습니다.

 

문제 제목: 야근 지수 (3단계)

문제 사이트: https://school.programmers.co.kr/learn/courses/30/lessons/12927?language=python3

프로그래머스

 

문제 설명

 

---

 

나의 풀이

 

import heapq

def solution(n, works):
    if n >= sum(works):
        return 0
    
    works = [-w for w in works]
    heapq.heapify(works)
    
    for _ in range(n):
        i = heapq.heappop(works)
        i += 1
        heapq.heappush(works, i)
    
    answer = sum([w ** 2 for w in works])
    return answer

 

n이 works 리스트에 있는 모든 원소의 합보다 크면 잔업이 없으므로 0를 반환하면 된다.

 

나머지는 works의 원소들을 음수로 바꾸고 자동으로 정렬해주는 heapq의 자료구조를 이용하여 작은 수부터 하나씩 빼서 1를 더해주고 다시 heapq에 다시 넣어준다.

 

이후 배열에 있는 원소의 제곱을 더해주고 answer에 대입한다.

---

※ 알아야 할 것

 

heapq.heappush, heapq.heappop을 통해서 원소를 제거하면서 heap의 자료구조가 heap인줄 알았다. 근데 풀다보니 그냥 heap(자료이름)의 자료 형태는 list였다. 그래서 list에서 사용할 수 있는 메소드들을 모두 사용할 수 있다는 사실을 알게되었다.

 

자세한건 이 문제에서 자세히 알 수 있다.