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/
'코딩테스트 > programmers (python)' 카테고리의 다른 글
Programmers / 3단계 / 등굣길 / python / 동적계획법(Dynamic Programming) (0) | 2024.03.28 |
---|---|
Programmers / 2단계 / 소수 찾기 / python / 완전 탐색 (0) | 2024.03.28 |
Programmers / 2단계 / [1차]프렌즈4블록 / python / 2018 KAKAO BLIND RECRUITMENT (1) | 2024.03.26 |
Programmers / 2단계 / 2 x n 타일링 / python (0) | 2024.03.25 |
Programmers / 3단계 / 최고의 집합 / python (0) | 2024.03.23 |