[프로그래머스] lv3. 야근 지수 풀이 (Python)

프로그래머스

접근한 방법

배열의 값들 중 가장 큰 값을 찾아서 1씩 제거해주어야 합니다. 이중 for문을 쓰면 시간 초과가 발생하므로 힙 구조를 이용해 제거해주면 됩니다. 주어진 works를 맥스 힙으로 구성한 후 n만큼 반복하여 매 순회 시 힙에서 값을 1제거한 후 다시 넣어주는 방식을 사용합니다. 자료를 넣고 빼는 과정은 각각 logn시간이 걸리므로 제한 시간 내에 풀이할 수 있습니다.

def solution(n, works):
    import heapq as hq
    q = []
    for i in works:
        hq.heappush(q, -i)
    for i in range(n):
        now = -hq.heappop(q)
        if not now:
            return 0
        hq.heappush(q, -(now - 1))
    return sum(map(lambda x: x ** 2, q))