[🥈2 / 백준 11279 / 파이썬] 최대 힙 (heapq 사용 x)

 

11279번: 최대 힙

문제

널리 잘 알려진 자료구조 중 최대 힙이 있다. 최대 힙을 이용하여 다음과 같은 연산을 지원하는 프로그램을 작성하시오.

1. 배열에 자연수 x를 넣는다.
2. 배열에서 가장 큰 값을 출력하고, 그 값을 배열에서 제거한다.

프로그램은 처음에 비어있는 배열에서 시작하게 된다.

입력

첫째 줄에 연산의 개수 N(1 ≤ N ≤ 100,000)이 주어진다. 다음 N개의 줄에는 연산에 대한 정보를 나타내는 정수 x가 주어진다. 만약 x가 자연수라면 배열에 x라는 값을 넣는(추가하는) 연산이고, x가 0이라면 배열에서 가장 큰 값을 출력하고 그 값을 배열에서 제거하는 경우이다. 입력되는 자연수는 231보다 작다.

출력

입력에서 0이 주어진 회수만큼 답을 출력한다. 만약 배열이 비어 있는 경우인데 가장 큰 값을 출력하라고 한 경우에는 0을 출력하면 된다.

 

예제 입력  13 0 1 2 0 0 3 2 1 0 0 0 0 0

예제 출력  0 2 1 3 2 1 0 0

 

풀이

 

힙 자료구조를 이용해 무작위 입력을 받는 배열의 최댓값만 계속하여 출력하는 문제이다.

파이썬의 경우 heapq 모듈을 사용하여 손쉽게 힙을 구현할 수 있지만 공부 목적을 위해 직접 구현해보았다.

 

우선, 힙은 완전 이진 트리이다.

때문에 자식 노드가 무조건 2개 이하이고, 왼쪽 자식 노드가 없고 오른쪽 자식 노드만 있는 경우가 없다. 순서대로 채워지기 때문이다.

 

최대 힙인 경우 부모 노드가 자식 노드보다 무조건 크고, 최소 힙일 경우 부모 노드가 자식 노드보다 무조건 작을 뿐이다.

 

트리는 아래 그림과 같이 구성되고, 1차원 배열에 기록될 때에는 부모의 현재 인덱스가 n이라면 왼쪽 자식 노드는 n*2, 오른쪽 자식 노드는 n*2 +1에 위치된다.

 

출처 - 위키피디아

그림에서는 인덱스가 0부터 시작하지만, 실제 구현을 하기 위해서는 0을 비워두고 인덱스 1부터 자료를 채워 넣는다.

부모 노드가 n이라면 자식 노드는 각각 n*2, n*2+1 로 구성되므로 인덱스가 0일 경우 조금 곤란해지기 때문이다.

 

 

힙에 자료를 push할 경우에는 맨 끝 노드부터 채워지게 된다. 완전 이진 트리이니 당연한 수순이다.

맨 끝에 push한 뒤, 부모 노드를 검사하며 부모 노드보다 크다면 위치를 바꾸며 거슬러 올라가는 과정이 필요한데, 이를 heapify라 한다.

 

반대로 pop을 할 경우, 무조건 뿌리 노드를 추출한다. 당연히 뿌리 노드에 최소, 혹은 최댓값이 들어있기 때문이다.

그리고 빈 뿌리 노드에 맨 마지막에 넣었던 노드를 가져온다.

그리고 아까 위에서 했던 heapify의 방향만 바꿔서 실행한다.

 

단, 부모를 거슬러 올라가는 과정과 달리 자식을 거쳐 내려가는 heapify는 두 자식 중 더 큰 것을 선별하여 내려가야 한다.

 

 

import sys
input = sys.stdin.readline

heap = [0]
capacity = 0

def push(heap, n):
    global capacity
    # heap의 자리가 부족할 경우 배열을 2배로 늘려나간다.
    if len(heap)-1 == capacity:
        heap = heap + [0]*len(heap)
    
    capacity += 1
    heap[capacity] = n
    
    # above heapify
    tn = capacity
    while(tn > 1):
        if heap[tn] > heap[tn//2]:
            temp = heap[tn]
            heap[tn] = heap[tn//2]
            heap[tn//2] = temp
            tn//=2
        else:
            break
    return heap

def pop(heap):
    global capacity
    if capacity == 0:
        return 0
    p = heap[1]
    heap[1] = heap[capacity]
    heap[capacity] = 0
    
    # below heapify
    tn = 1
    while(capacity > tn*2):
        if heap[tn*2] == 0 and heap[tn*2+1] == 0:
            break
        if heap[tn] < max(heap[tn*2], heap[tn*2+1]):
            temp = heap[tn]
            if heap[tn*2] > heap[tn*2+1]:
                heap[tn] = heap[tn*2]
                heap[tn*2] = temp
                tn *= 2
            else:
                heap[tn] = heap[tn*2+1]
                heap[tn*2+1] = temp
                tn = tn*2+1
        else:
            break
    capacity -= 1
    return p
    

for i in range(int(input())):
    x = int(input())
    if x:
        heap = push(heap,x)
    else:
        print(pop(heap))

 

참고 : https://hongjw1938.tistory.com/22