[알고리즘] 프로그래머스 야근 지수 파이썬 | 힙

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

문제

나의 풀이

제곱의 합을 최소화하기 위해서는 총합을 최소화 해야했다.

따라서 works 리스트 내의 최댓값을 찾아 1씩 감소시키는 것을 n번 반복했다.
하지만 max()와 index() 함수를 사용하는 과정에서 O(N)의 시간 복잡도가 소요되어
전체적으로 O(N^2)의 시간 복잡도가 발생해 효율성 테스트에서 시간이 초과되었다.

def solution(n, works):
    for _ in range(n): # O(N)
        if max(works) == 0:
            break
        works[works.index(max(works))] -= 1 # O(N)
        
    return sum(list(map((lambda x : x**2), works)))

max heap을 사용하라는 힌트를 얻어 다시 코드를 작성했다.

from heapq import heapify, heapreplace
def solution(n, works):
    answer = 0
    heap = [(-work, work) for work in works]
    heapify(heap)
    for _ in range(n): # O(N)
        if heap[0][0] == 0:
            break
        heapreplace(heap, (heap[0][0]+1, heap[0][1]-1)) # O(logN)
    return sum(list(map(lambda x:x[0]**2, heap)))

힙을 사용함으로써 시간 복잡도를 O(NlogN)로 줄였기에 효율성 테스트에서도 통과된 듯 하다.

다른 풀이

다른 풀이를 확인하고 생각해보니 튜플이 굳이 필요하지 않았다.
또한 n과 남은 작업량의 합 중 최솟값만큼의 반복문을 돌림으로써 if문이 불필요해졌다.

from heapq import heapify, heappush, heappop
def solution(n, works):
    heapify(works := [-i for i in works])
    for i in range(min(n, abs(sum(works)))):
        heappush(works, heappop(works)+1)
    return sum([i*i for i in works])

얻어가기

A := B

더 큰 표현식의 일부로 변수에 값을 대입하는 새로운 문법이다.