백준 1927번과 파이썬 힙

백준 1927번 문제는:

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

1927번: 최소 힙

전공에서 자료구조를 배울 때 힙을 C로 구현해서 코드도 길고... 짜다보니 집중력이 떨어져서 실수를 많이했던 기억이 있다. 

그런 힙을 파이썬으로 구현하려고 생각해보니 뭔가 heap이라는게 존재할 것 같아서 찾아보니 실제로

import heapq

로 힙을 사용할 수 있게 파이썬에 모든게 준비되어 있다는 걸 알게되었다. 야호

 

힙 함수들

1. heapq.heappush

heap = [] #비어 있는 리스트 생성

#heapq.heappush(리스트, 넣고 싶은 숫자)
heapq.heappush(heap, 0)
heapq.heappush(heap, 1)

print(heap)

=====================================
<코드 run>

[0, 1]

조금 독특하게도 heappush는 앞에 내가 먼저 생성해둔 리스트에 값을 넣고, 그 값을 넣으면 알아서 최소 힙의 기준에 맞게 

앞에서부터 값이 가장 작은 원소를 정렬해서 보여준다. 

2. heap1.heappop

print(heapq.heappop(heap))

===============================
<코드 run>
0

heappop은 최소 힙의 root, 즉 가장 작은 값을 return한다. 그래서 새로운 값에 저장해서 print하거나 그냥 print해도 값을 얻을 수 있고, 그 값은 내가 만든 heap이라는 list에서 사라진다.

 

문제 해결(1927번)

어차피 백준에도 소스 공개를 해두는 편이라 부족하지만.. 소스를 공개해보자면,

 

from sys import stdin
import heapq

def inputZero():
    try:
        print(heapq.heappop(heap))
    except:
        print(0)

def inputElse(num):
    heapq.heappush(heap,num)

heap = []

N = int(stdin.readline())
for _ in range(0,N):
    num = int(stdin.readline())
    if num == 0:
        inputZero()
    else:
        inputElse(num)

이렇게 해결했다. N을 입력 받아 for문을 돌리며 한 번씩 입력 받고 바로 결과를 출력하는 형식으로 구성했다.

(stdin.readline()을 사용하는 이유는 백준에서 시간 초과가 자주 발생하기 때문..)

 

만약 0을 입력 받으면

def inputZero():
    try:
        print(heapq.heappop(heap))
    except:
        print(0)

inputZero()라는 함수를 불러오는데, 이때 try except 를 사용했다. 만약 heapop(heap)에서 인덱스가 비어있으면 indexError가 발생하는데, 그때 바로 except로 넘어가서 0을 print해주는 식이다.

 

def inputElse(num):
    heapq.heappush(heap,num)

그리고 일반적으로 추가할 숫자를 입력 받으면 InputElse(num)이라는 함수로 push해준다.

 

사실 C에 비하면.. 자료구조 자체를 완벽히 알고 있어서 코드를 짤 수 있다기보다는 그냥 개념만 알고 뭘 Import 할지 알고 있다면 금방 해결할 수 있는 문제처럼 느껴진다. 다만 시간 복잡도가 늘 아쉽고, C로 새로 짜면서 이해 하는게 중요할 듯 하다.