[백준 11279] 최대 힙 with Python

문제 링크

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

11279번: 최대 힙

풀이

import sys

class Heap:
	def __init__(self):
		self.arr = []
		self.output = []
	
	def heapify_down(self, k, n):
		while 2*k+1 < n:
			L, R = 2*k+1, 2*k+2
			if k < n and self.arr[L] > self.arr[k]:
				m = L
			else:
				m = k
			if R < n and self.arr[R] > self.arr[m]:
				m = R
			if m != k:
				self.arr[m], self.arr[k] = self.arr[k], self.arr[m]
				k = m
			else:
				break

	def heapify_up(self, k):
		parent = (k-1)//2
		while k > 0 and self.arr[parent] < self.arr[k]:
			self.arr[parent], self.arr[k] = self.arr[k], self.arr[parent]
			k = parent
			parent = (k-1)//2

	def insert(self, key):
		self.arr.append(key)
		self.heapify_up(len(self.arr)-1)

	def delete_max(self):
		if len(self.arr) == 0:
			self.output.append(0)
			return
		self.arr[0], self.arr[len(self.arr)-1] = self.arr[len(self.arr)-1], self.arr[0]
		self.output.append(self.arr.pop())
		self.heapify_down(0, len(self.arr))

	def print(self):
		if self.output:
			print("\n".join(map(str, self.output)))

N = int(sys.stdin.readline())
heap = Heap()
for _ in range(N):
	x = int(sys.stdin.readline())
	if x: heap.insert(x)
	else: heap.delete_max()
heap.print()

힙 자료구조를 직접 구현하여 해결하였다.

 

Python의 heapq 모듈을 사용하면 heap을 직접 구현하지 않고도 쉽게 해결할 수 있다.

import sys
import heapq

N = int(sys.stdin.readline())
heap = []
output = []
for _ in range(N):
	x = int(sys.stdin.readline())
	if x: heapq.heappush(heap, -x)
	else: output.append(-heapq.heappop(heap)) if heap else output.append(0)
print("\n".join(map(str, output)))

단, heapq는 최대힙(Max Heap)이 아닌 최소힙(Min Heap)을 구현한 모듈이기 때문에 이 문제를 풀 때에는 입력값과 출력값을 모두 부호를 바꿔줘야 한다.