[Python] 백준 11279 최대 힙

문제

최대 힙을 이용하여 다음과 같은 연산을 지원하는 프로그램을 작성

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

프로그램은 처음에 비어있는 배열에서 시작

 

 

 

입력

• 첫째 줄: 연산의 개수 N(1 ≤ N ≤ 100,000)
• N개의 줄 : 연산에 대한 정보를 나타내는 정수 x

 

 

 

출력

입력에서 0이 주어진 횟수만큼 답을 출력( 입력이 0이면 가장 큰 값 출력, 자연수라면 heap에 해당 원소 추가)만약 배열이 비어 있는 경우인데 가장 큰 값을 출력하라고 한 경우에는 0을 출력

 

 

https://www.acmicpc.net/problem/11279

11279번: 최대 힙

 

---

★Keypoint : 힙, sys

 

heapq 라이브러리를 활용하여 계산 수행하기

 

 

<Try CODE>

import heapq

class Heap:
    def __init__(self, size):
        self.size = size
        self.heap = [None] * size
        self.count = 0
    def isEmpty(self):
        return self.count == 0 #
    def isFull(self):
        return self.size -1 == self.count

    def add_heap(self, item):
        if self.isFull():
            self.remove_max()
        self.count += 1
        i = self.count
        while i != 1 and item > self.heap[i // 2]:
            self.heap[i] = self.heap[i // 2]
            i //= 2
        self.heap[i] = item

    def remove_max(self):
        if self.isEmpty(): return None
        max_val = self.heap[1] #가장 큰 값 뿌리
        temp = self.heap[self.count] #가장 끝 원소 임시 저장
        self.count -= 1 #가장 큰 값이 빠질 예정이므로 전체 원소 크기 하나 줄이기
        parent = 1 #힙을 재구성하기 위해서 parent 1로 두기
        while parent * 2 <= self.count: #현재 부모의 왼쪽 자식이 힙 범위 내에 있는 한
            child = parent * 2
            if child < self.count and self.heap[child] < self.heap[child + 1]:
                child += 1 #오른쪽 자식이 있는지 확인하고 (child < self.count) 오른쪽 자식이 왼쪽 자식보다 큰지 확인 참일 경우 child를 오른쪽 자식을 나타내도록 증가
            if temp >= self.heap[child]: #temp에 저장된 값 (힙의 마지막 원소)이 현재 고려 중인 자식의 값보다 크거나 같은지 확인
                break  #참일 경우 힙 속성이 유지되므로 루프를 종료
            self.heap[parent] = self.heap[child]
            parent = child #다음 루프 반복 준비
        self.heap[parent] = temp

        return max_val

N = int(input())
heap = Heap(N + 1)  # 인덱스 0은 dummy 값

for _ in range(N):
    x = int(input())
    if x == 0:
        if heap.isEmpty():
            print(0)
        else:
            print(heap.remove_max())
    else:
        heap.add_heap(x)

 

처음에는 이 문제를 heapq 라이브러리를 활용하지 않고 구현하려 했다.

오랜만에 class도 사용하면서 자료구조 시간에 배웠던 알고리즘으로 제출했지만, 시간 초과라는 결과가 나왔다.

 

 

 

<최종 CODE>

import heapq
import sys

class MaxHeap:
    def __init__(self):
        self.heap = []
    
    def push(self, value):
        heapq.heappush(self.heap, -value)
    
    def pop(self):
        if self.heap:
            return -heapq.heappop(self.heap)
        return 0

N = int(sys.stdin.readline().strip())
heap = MaxHeap()

result = []
for _ in range(N):
    x = int(sys.stdin.readline().strip())
    if x == 0:
        result.append(heap.pop())
    else:
        heap.push(x)

sys.stdout.write('\n'.join(map(str, result)))

 

또한, sys를 import 하지 않는다면 시간 초과 결과를 피하기 어렵다!

해당 코드에서 ' - ' 를 사용한 것을 볼 수 있는데, 사용한 이유는 Python의 heapq 라이브러리가 기본적으로 최소 힙을 지원하기 때문이다!

 

코드에서 -item을 저장하고, 값을 사용할 때는 -heapq.heappop(self.heap)을 통해 다시 원래의 값으로 되돌리는 것! 

 

이 노드에서 가장 상단(root)에 위치한 노드를 추출하고 음수 부호를 다시 입혀서 양수로 만든 다음 출력하는 것이다!