알고리즘/Heap 힙

[프로그래머스] 더 맵게 - 파이썬(python)

아뵹젼 2022. 9. 3.

 

 

나의 풀이

from heapq import heapify, heappush, heappop
def solution(scoville, K):
    answer = 0
    heapify(scoville)
    
    while True :
        if scoville[0] >= K :
            break
        elif len(scoville) < 2:
            answer = -1
            break
        else :
            heappush(scoville, heappop(scoville) + heappop(scoville)*2)
        answer += 1
    return answer

처음에는 sort() 를 이용한 정렬로 풀었는데, 효율성 테스트에서 가차없이 실패하였다.

따라서 heapq를 이용해서 같은 로직으로 다시 풀어보았다.

scoville 리스트를 heapify() 를 이용하여 heap 로 변경하였다.

 

반복문을 통해 가장 작은 원소인 0번째 원소가 K 이상일 때까지 검사한다.

만약 0번째 원소가 K보다 작은데 scoville 에 남아있는 원소가 1개 혹은 0개라면, 이는 스코빌 지수를 K이상으로 만들 수 없음을 의미하므로 -1 을 리턴한다.

원소가 2개 이상이라면, 첫번째 원소 + 두번째로 작은원소*2 를 heappush 한다.

 

 

다른 풀이

from heapq import heapify, heappush, heappop
def solution(scoville, K):
    heapify(scoville)
    for i in range(1000000):
        try:
            heappush(scoville, heappop(scoville)+(heappop(scoville)*2))
            if scoville[0] >= K: return i+1
        except:
            return -1

예외처리로 깔끔한 코드..!

만약 scoville 에 남아있는 원소가 1개 혹은 0개라면 try 문에서 인덱스 에러가 발생할 것이다 (heappop() 때문에)

따라서 except 로 가서 -1 을 리턴하게 된다.

 

파이썬 heapq 모듈

heapq 모듈은 이진 트리기반의 최소 힙(min heap) 자료구조를 제공한다.

즉, 항상 가장 작은 값이 인덱스 0 번에 위치하므로, 최댓값 혹은 최솟값을 찾기에 적합하다.

최대 힙을 사용하고 싶다면? 편법으로 heapq.heappush(리스트, -x) 를 넣으면 된다.

그리고 -heapq.heappop(리스트) 로 음수 값을 다시 양수로 되돌리면 최대 힙처럼 사용할 수 있다.

 

이미 원소가 들어있는 리스트 힙으로 만들려면 heapq 모듈의 heapify()라는 함수를 사용하면 된다.

heap = [4, 1, 7, 3, 8, 5]
heapify(heap)

 

heapq.heappushpop(heap, item)

힙에 item을 푸시한 다음, heap에서 가장 작은 항목을 팝하고 반환한다.

heappush() 를 한 다음 heappop() 를 별도로 호출하는 것보다 더 효율적이다!!!!

따라서 heappush(scoville, heappop(scoville)+(heappop(scoville)*2)) 를 다음과 같이 사용할 수 있다.

s = heapq.heappop(scoville)
f = heapq.heappushpop(scoville, f + s + s)
  • heapq.nlargest(n, iterable, key=None)
    • n개의 가장 큰 요소로 구성된 리스트 반환
    • 따라서 반환된 리스트의 n번째 요소는 heapq 에서 n번째로 큰 원소일 것이다.
  • heapq.nsmallest(n, iterable, key=None)
    • n개의 가장 작은 요소로 구성된 리스트 반환
    • 따라서 반환된 리스트의 n번째 요소는 heapq 에서 n번째로 작은 원소일 것이다.

 

댓글