최대 힙
문제
널리 잘 알려진 자료구조 중 최대 힙이 있다. 최대 힙을 이용하여 다음과 같은 연산을 지원하는 프로그램을 작성하시오.
- 배열에 자연수 x를 넣는다.
- 배열에서 가장 큰 값을 출력하고, 그 값을 배열에서 제거한다.
프로그램은 처음에 비어있는 배열에서 시작하게 된다.
입력
첫째 줄에 연산의 개수 N(1 ≤ N ≤ 100,000)이 주어진다. 다음 N개의 줄에는 연산에 대한 정보를 나타내는 정수 x가 주어진다. 만약 x가 자연수라면 배열에 x라는 값을 넣는(추가하는) 연산이고, x가 0이라면 배열에서 가장 큰 값을 출력하고 그 값을 배열에서 제거하는 경우이다. 입력되는 자연수는 2^31보다 작다.
출력
입력에서 0이 주어진 횟수만큼 답을 출력한다. 만약 배열이 비어 있는 경우인데 가장 큰 값을 출력하라고 한 경우에는 0을 출력하면 된다.
-
이 문제는 자료 구조 heap에서 최대 힙을 사용할 수 있는가?
-
커다란 입력을 받을때 이것을 시간 초과가 뜨지 않게 받을 수 있는가?
크게 2가지 개념을 묻는 문제다.
heap을 사용할 수 있는가?
문제에서 대놓고 heap을 사용하라고 나와있으니 잠깐 heap에 대해서 설명하자면 다음과 같다.
힙은 특정한 규칙을 가진 트리 기반의 자료구조이다.
이것은 우선순위 큐를 구현하는데 사용된다.
결과적으로 최댓값과 최솟값을 찾는 연산을 빠르게 하기 위해 사용한다.
-> 완전이진트리 형태
파이썬에서 heap 모듈 함수들
-
heapq.heappush(heap, item)
이 함수는 아이템을 힙에 추가한다. -
heapq.heappop(heap)
힙에서 가장 작은 아이템을 뺴고 그 값을 반한다. 이 함수 자체가 값을 반환하기 때문에
print를 쓰면 값을 출력할 수 있다. -
heapq.heappushpop(heap, item)
아이템을 힙에 추가한 다음, 가장 작은 아이템을 빼고 그 값을 반환한다. -
heapq.heapify(x)
리스트 x를 즉각적으로 heap으로 변환한다. (O(N))
명심할 것.
위에 나온 heap 함수들은 최소 힙일때 사용한다. 파이썬의 heapq 모듈은 처음엔 최소 힙으로 사용할 수 있도록 되어있다. 반대로 최대힙을 사용하려면 위에 있는 함수를 조금 변형해야 한다.
부호만 바꾸는법
heap에 push를 할때 원래는
- heapq.heappush(heap, item)
이렇게 했지만 아래 처럼 부호를 바꿔서 넣어야 한다. (최대 힙으로 사용하고자 할 경우)
- heapq.heappush(heap, -item)
어차피 나중에 출력을 할때 부호를 다시 바꿔서 출력하면 되기 때문에 최소 힙으로 사용하도록 세팅되어있는 push기능에 -를 함으로써 우선순위를 설정한 다음 결과를 출력할때 부호를 다시 바꿔 주면 된다. 이 문제는 이 방법으로 풀었다.
부호와 튜플을 이용하는법.
튜플을 이용하는 방법도 있다.
push를 할때 아래 처럼 한다.
max_heap = []
numbers = [4,1,7,3,8,5]
for item in numbers:
heapq.heappush(max_heap, (-item, item)) 이런식으로 (-item, item) 튜플을 max_heap에 넣으면 각 튜플의 0번째에 음수로 되어있는 것들이 최소 힙처럼 우선순위 큐가 된다.
나중에 최대 값을 빼서 출력할때는 아래처럼 하면 된다.
heapq.heappop(max_heap)[1]
튜플의 1번째가 실제 값이므로 이렇게 하면 된다.
풀이
import heapq
n = int(input())
max_heap = []
for _ in range(n):
number = int(input())
if number > 0:
heapq.heappush(max_heap, -number)
else:
if not max_heap:
print(0)
else:
print(-heapq.heappop(max_heap))
하지만 이렇게 풀면 시간초과로 틀린다.
이 문제는 바로 2번째 질문을 하고 있다. 문제에서는 입력되는 자연수가 2^31보다 작다고 되어있다. 입력되는 수가 거대하므로 이것을 int(input())으로 받으면 시간초과가 난다. 따라서 이것을 아래 처럼 받아야 시간초과가 나지 않는다.
import sys
import heapq
n = int(sys.stdin.readline())
max_heap = []
for _ in range(n):
number = int(sys.stdin.readline())
if number > 0:
heapq.heappush(max_heap, -number)
else:
if not max_heap:
print(0)
else:
print(-heapq.heappop(max_heap))
sys.stdin.readline()로 하면 시간초과를 낼 수가 없다!
