문제링크: https://school.programmers.co.kr/learn/courses/30/lessons/12927
문제를 풀다가 오랜만에 heap문제가 나와서 포스팅해보겠습니다!
Problem
숫자열과 숫자 n 이 주어집니다. n을 숫자열에 적절히 분배하고 빼서 숫자열 원소의 제곱 합이 최소가 되게 하는 문제입니다.
말로 하려니 너무 어렵네요ㅋㅋㅋ
예를들어 [4, 3, 3]와 4가 있다면, 첫 원소에서 2를, 나머지 원소에서 1씩 빼서, 결과적으로 [2,2,2]을 만들고, 각 원소의 제곱 합, 즉 12를 반환하면 되는 문제입니다.
Approach & Solution
처음에는 당연히 단순 리스트로 접근했습니다. 가장 큰 수를 찾고, 그 수에 1을 빼는 방식으로 접근했습니다. while 종료 조건은 최대값이 0이거나, n이 0이 될 때 입니다.
def solution(n, works):
answer = 0
idx = 0
while n > 0:
if works[works.index(max(works))] > 0:
works[works.index(max(works))] -= 1
n -= 1
idx += 1
idx %= len(works)
for i in works:
answer = answer + i**2
return answer근데 정확성은 다 맞는데, 효율성에서 점수가 안나오더라구요. 이미 works[works.index(max(works))] 에서 O(n^2)을 넘어버립니다. 그래서 더 빠른 방법이 뭐지하고 힌트를 살짝 봤는데, 힙을 사용하더라구요.
import heapq as q
def solution(n, works):
heap = []
for i in works:
q.heappush(heap, -i)
answer = 0
idx = 0
while n > 0:
work = q.heappop(heap)
if work == 0:
break
q.heappush(heap,work+1)
n -= 1
idx += 1
idx %= len(works)
print(heap)
for i in heap:
answer = answer + i**2
return answer힙은 보통 최소힙을 말합니다. 파이썬도 최소힙밖에 지원하지 않습니다. 하지만 여기서는 최대값을 찾는 문제이므로, 힙에 넣는 원소들을 음수로 바꿔서 사용하면 최대힙과 동치가 되는 방법으로 접근했습니다.
그리고 두번째 팁으로 힙에 정보를 갱신하려면 heap[0] 을 업데이트하고 다시 heapify 해주는것이 아닌, heappop()으로 최대값을 얻고, 다시 heappush() 하는 방법으로 접근해야 더 빠르게 됩니다. 위의 풀이에서도 work에 최대값을 받고, 다시 q.heappush(heap,work+1) 하는 방식으로 접근해 봤습니다.
Conclusion
오랜만에 힙 문제를 풀어보았습니다. 역시 꾸준히 풀어야 되는것 같습니다. 비교적 간단한 힙 문제에서 여러 개념을 리마인드 할 수 있는 기회였던것 같습니다.
