Programmers / 3단계 / 야근 지수 / python

 

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

 

 

나의 풀이

def solution(n, works):
    answer = 0
    if sum(works) <= n:
        return 0
    
    for i in range(n):
        works[works.index(max(works))] -= 1
    
    for i in works:
        answer += i**2
    
    return answer

 

인덱스를 사용해 works에 있는 가장 큰 수에서 -1을 해 정답을 찾았지만 시간 초과로 실패하였다.

 

 

시간을 단축시키기 위해 heapq를 이용하여 풀었다.

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

 

 

모범 답안

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 = [1, 2, 3]
if (n := len(a)) >5:
	print('hello')

위 예시처럼 대입 표현식을 사용하면 n을 조건문 안에서 선언하고 값을 대입할 수 있다.

 

 

 

참조

https://int-i.github.io/python/2020-05-29/python-walrus-operator/

Python 3.8 기여운 바다코끼리를 드리겟슴미다 := - 인하대학교 인트아이