2 분 소요

최소 힙

solved_ac[Class3] 최소 힙

문제

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

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

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

입력

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

출력

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

예제 입력 1

9
0
12345678
1
2
0
0
0
0
32

예제 출력 1

0
1
2
12345678
0

문제 해석

문제에 나와있는데로 힙을 사용하는 문제이다. 파이썬에서는 최소힙만 지원을 해주기 때문에 만약 최대힙을 구하고 싶다면 음수로 집어넣어서 뺀 후 다시 양수로 바꿔주면 된다.

힙(heap)

image

최댓값과 최솟값을 빠르게 찾기 위해 고안된 자료구조

  • 각 노드의 key 값이 해당 노드의 자식 노드의 key 값보다 작지 않거나 크지 않은 완전 이진 트리
  • key 값의 대소 관계는 부모 - 자식 노드 사이 간에만 성립하며 형제 노드 사이에는 영향을 미치지 않는다.
  • 자식노드의 최대 개수는 힙의 종류에 따라 다르지만 이진 트리에서는 최대 2개(완전 이진 트리를 사용한다고 가정)이다.
  • i 번째 노드의 자식 노드가 2개인데 왼쪽 자식 노드는 2i, 오른쪽 자식 노드는 2i + 1 이고, 부모 노드는 i / 2가 된다.

힙(heap)의 시간 복잡도

O(logN)

힙(heap)의 연산 과정

삽입 연산(insertion)

  • 새로운 원소가 삽입되었을 때 O(logN)의 시간 복잡도로 힙 성질을 유지하도록 할 수 있다.

image

  1. 삽입하고자 하는 값을 트리의 가장 마지막 원소에 추가한다.
  2. 부모 노드와의 대소 관계를 비교하면서 만족할 때까지 자리 교환을 반복한다.

삭제 연산(deletion)

image

image

  1. 힙에서는 루트 노드만 삭제가 가능하므로 루트 노드를 제거한다.
  2. 가장 마지막 노드를 루트로 이동시킨다.
  3. 자식 노드와 비교하여 더 작은 자식 노드로 조건이 만족할 때까지 이동시킨다.

힙큐 모듈(heapq module)

모듈 임포트

import heapq

노드 추가

  • heappush 메소드 이용
heap = []
heapq.heappush(heap, 1)

노드 삭제

  • heappop 메소드 이용
  • 가장 작은 원소를 꺼내어 리턴, 자동적으로 그 다음으로 작은 원소가 루트 노드가 됨
ans = heapq.heappop(heap)
print(ans)

# 만약 꺼내지 않고 출력만 하고 싶다면?
# heappop 메소드를 이용하지 않고 인덱스로 접근
print(heap[0])

주의

  • 만약 n번째로 작은 원소를 얻고 싶다면 n-1 인덱스로 접근을 하면 안된다.
  • 이진 트리의 루트 노드는 가장 작거나 가장 큰게 확실한데, 그 자식 노드들은 순서대로 정렬되어 있는 것이 아니기 때문이다.
  • 그래서 n번째로 작은 원소를 얻고 싶다면 n-1번 heappop(heap)을 해줘야 한다.
  • 그래야 루트 노드가 n번째로 작은 원소가 된다.

풀이(최소 힙)

  • 입력 받은 수가 0이라면?
    • heap의 길이가 0이면 비어있다는 얘기이기 때문에 0을 출력해준다.
    • heap이 안비어있다면 heappop()을 해서 최소값을 출력해준다.
  • 0이 아니라면?
    • 입력 받은 수를 heappush()를 이용하여 집어넣어준다.
import sys
import heapq

N = int(sys.stdin.readline())

heap = []

for i in range(N):
    a = int(sys.stdin.readline())
    if a == 0:
        if len(heap) == 0:
            print(0)
        else:
            print(heapq.heappop(heap))
    else:
        heapq.heappush(heap, a)

풀이(최대 힙)

파이썬에서 제공해주는 heapq는 최소힙이다. 만약 최대 힙 문제를 풀고 싶다면?

  • 입력 받은 수가 0이라면?
    • heap의 길이가 0이면 비어있다는 얘기이기 때문에 0을 출력해준다.
    • heap이 안비어있다면 heappop()을 해서 최소값을 뺀다.
    • 음수로 집어넣어줬기 때문에 다시 -를 붙여서 양수로 출력해준다.
  • 0이 아니라면?
    • 입력 받은 수를 heappush()를 이용하여 집어넣어준다.
    • 최소 힙으로 되어있기 때문에 -값으로 집어넣어줘야 한다.
import sys
import heapq

N = int(sys.stdin.readline())

heap = []

for i in range(N):
    a = int(sys.stdin.readline())
    if a == 0:
        if len(heap) == 0:
            print(0)
        else:
            print(-heapq.heappop(heap))
    else:
        heapq.heappush(heap, -a)

댓글남기기