파이썬 힙 함수 (프로그래머스 <명예의 전당(1)> , <더 맵게> , <야근 지수> )

파이썬 힙 자료구조

파이썬 heapq 모듈은 heapq (priority queue) 알고리즘을 제공한다.

 

모든 부모 노드는 그의 자식 노드보다 값이 작거나 큰 이진트리(binary tree) 구조인데, 내부적으로는 인덱스 0에서 시작해 k번째 원소가 항상 자식 원소들(2k+1, 2k+2) 보다 작거나 같은 최소 힙의 형태로 정렬된다.   

 

heapq는 내장 모듈로 별도의 설치 작업 없이 바로 사용할 수 있다.

 

힙 함수 활용하기

• heapq.heappush(heap, item) : item을 heap에 추가
• heapq.heappop(heap) : heap에서 가장 작은 원소를 pop & 리턴. 비어 있는 경우 IndexError가 호출됨. 
• heapq.heapify(x) : 리스트 x를 즉각적으로 heap으로 변환함 (in linear time, O(N) )

힙 생성 & 원소 추가

heapq 모듈은 리스트를 최소 힙처럼 다룰 수 있도록 하기 때문에, 빈 리스트를 생성한 후 heapq의 함수를 호출할 때마다 리스트를  인자에 넘겨야 한다.

 

아래 코드는 heap이라는 빈 리스트를 생성하고 50, 10, 20을 각각 추가하는 예시이다.

import heapq

heap = []
heapq.heappush(heap, 50)
heapq.heappush(heap, 10)
heapq.heappush(heap, 20)

print(heap)

이미 생성해둔 리스트가 있다면 heapify 함수를 통해 즉각적으로 힙 자료형으로 변환할 수 있다.

heap2 = [50 ,10, 20]
heapq.heapify(heap2)

print(heap2)

위의 예제에서 힙에서 가장 작은 원소인 20을 가져오고난 후에도 heap은 유지되는 것을 확인할 수 있다. 

최대 힙 만들기

파이썬의 heapq 모듈은 최소 힙으로 구현되어 있기 때문에 최대 힙 구현을 위해서는 트릭이 필요하다.

 

IDEA: y = -x 변환을 하면 최솟값 정렬이 최댓값 정렬로 바뀐다.

 

힙에 원소를 추가할 때 (-item, item)의 튜플 형태로 넣어주면 튜플의 첫 번째 원소를 우선순위로 힙을 구성하게 된다.  이때 원소 값의 부호를 바꿨기 때문에, 최소 힙으로 구현된 heapq 모듈을 최대 힙 구현에 활용하게 되는 것이다.

 

아래의 예시는 리스트 heap_items에 있는 원소들을 max_heap이라는 최대 힙 자료구조로 만드는 코드이다.

heap_items = [1,3,5,7,9]

max_heap = []
for item in heap_items:
  heapq.heappush(max_heap, (-item, item))

print(max_heap)

 

heappush 함수를 통해 item을 힙에 push 할 때 (-item, item)의 튜플의 형태로 넣은 것을 확인할 수 있다.

 

그 결과 heappop을 사용하게 되면 힙에 있는 최댓값이 반환되는 것을 확인할 수 있다. 이때 실제 원소 값은 튜플의 두 번째 자리에 저장되어 있으므로 [1] 인덱싱을 통해 접근하면 된다.

max_item = heapq.heappop(max_heap)[1]
print(max_item)

 

 

명예의 전당

import heapq

def solution(k,score):
    answer = []
    rank  = []
    for s in score:
        heapq.heappush(rank,s)
        
        if len(rank) > k:
            heapq.heappop(rank)
            
        answer.append(rank[0])
    return answer

----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------

더 맵게

import heapq
def solution(scoville, K):
    count = 0 
    heapq.heapify(scoville) 
    

    while len(scoville) > 1:
 
        first = heapq.heappop(scoville)
        

        if first < K:
            count += 1                           
            second = heapq.heappop(scoville)           
            heapq.heappush(scoville, (first + (second * 2)))
        else:
            break
        
    return -1 if scoville[0] < K else count

----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------

야근 지수

from heapq import heappush, heappop

def solution(n, works):
    pq = []
    
    for work in works:
        heappush(pq, (-work, work))
        
    while n > 0:
        idx, val = heappop(pq)
        val -= 1
        if val <= 0:
            val = 0
        heappush(pq, (-val, val))
        n -= 1
    
    return sum([p[1]**2 for p in pq])