[Python] 야근 지수

https://school.programmers.co.kr/learn/courses/30/lessons/12927

프로그래머스

• 첫 번째 풀이(시간 초과)

# 야근 지수
# 효율성 통과 못함...
def solution(n, works):
    answer = 0
    
    for _ in range(n):
        if sum(works)==0:
            return 0
        works[works.index(max(works))] -= 1
    
    for i in works:
        if i!=0:
            answer += i**2
    return answer

 

• 두 번째 풀이(통과)

# 다른 풀이
# max heap을 이용한 풀이

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)
    
    return sum([w ** 2 for w in works])

이제 슬슬... 자료구조도 생각하면서 문제를 풀어야겠다..