알고리즘(4)

정길규·2023년 5월 25일

백준 11279번(최대 힙)

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

오답

문제만 읽었을 때에는 단순히 배열에 0이 아닌 숫자가 입력값에 들어오면 집어넣고 0이 나오면 배열에서 재일 큰수를 빼면 될줄 알고 sort()를 사용해서 pop()을 이용해서 풀었다.
그러니 시간초과가 나오게 되었다.

import sys
n = int(input())
list = []
for _ in range(n):
    num = int(sys.stdin.readline())
    if num == 0:
        if len(list) == 0:
            print(0)
        else:
            list.sort()
            print(list.pop())
    else:
        list.append(num)

heaqp

heapq 모듈은 이진 트리(binary tree) 기반의 최소 힙(min heap) 자료구조를 제공.
min heap을 사용하면 원소들이 항상 정렬된 상태로 추가되고 삭제되며, min heap에서 가장 작은값은 언제나 인덱스 0, 즉, 이진 트리의 루트에 위치합니다.

heaqp 생성 및 heappush()

heapppush(리스트명, 추가할 수)

from heapq import heappush, heappop

heap = []  # heap 생성

heappush(heap, 4)
heappush(heap, 1)
heappush(heap, 7)
heappush(heap, 5)
heappush(heap, 3)
heappush(heap, 8)

print(heap)   # [1, 3, 7, 5, 4, 8]

리스트와 다르게 삽입되는 순서대로 정렬되지 않는다. 위의 설명 대로 2진트리 형태로 저장된다.

재일 작은 수가 root가 되고 그 밑의 수는 위의 수보다 큰 수들이 위치하게 된다.

heappop()

from heapq import heappop

heap = [1, 3, 7, 4, 5]

heappop(heap)

print(heap)  # [3, 4, 7, 5]

heappop(heap)

print(heap)  # [4, 5, 7]

heappop()을 사용할 때마다 가장 작은 수가 삭제되고 리스트가 재배열 된다.

최소값 삭제하지 않고 접근하기

리스트에서 사용하는 방법대로 heap[0]을 통해 최소값에 접근할 수 있다.
단, heap[1]은 두번째로 작은 수가 아닐수도 있다.

from heapq import heappush

heap = []

heappush(heap, 4)
heappush(heap, 1)
heappush(heap, 7)
heappush(heap, 5)

print(heap)     # [1, 4, 7, 5]
print(heap[0])  # [1]

기존 리스트를 힙으로 변환

from heapq import heapify

heap = [5, 3, 7, 8, 1]
heapify(heap)

print(heap)   # [1, 3, 7, 8, 5]

리스트를 힙으로 변환시 내용의 순서도 바뀌게 된다. 단 새로운 힙리스트를 만드는게 아니라 기존의 리스트를 변환시키는 것이기 때문에 기존의 리스트를 보존할려면 리스트를 복제하여 힙으로 변환 시켜야 된다.

from heapq import heapify

list = [5, 3, 7, 8, 1]
heap = list[:]  # 복제

heapify(heap)

print(list)   # [5, 3, 7, 8, 1]
print(heap)   # [1, 3, 7, 8, 5]

[응용] 최대 힙

heap 모듈은 기본적으로 최소 힙 기반으로 제공된다.

from heapq import heappush, heappop

nums = [4, 1, 7, 3, 8, 5]
heap = []

for num in nums:
	heappush(heap, (-num, num)) # [(-8, 8), (-7, 7), (-5, 5), (-1, 1), (-3, 3), (-4, 4)]
    
while heap:
	print(heappop(heap)[1]) # 8, 7, 5, 4, 3, 2, 1 순으로 pop됨.

heap이 최소 값이 0 인덱스에 오는 것을 이용해서 튜플형태로 값을 저장 함. 튜플의 0 인덱스 위치에 num 음수를 붙여서 최대값을 최소값으로 만들고 반복문으로 튜플의 1 인덱스의 값만 pop하여 출력한다.

[응용] n번째 최소값 / 최대값

from heapq import heappush, heappop, heapify

def nth_smallest(nums, n):
    heap = []
    # heapify를 이용해서 for문을 생략할 수 있다.
	for num in nums:
    	heappush(heap, num)
        
    nth_min = None
    for _ in range(n):
    	nth_min = heappop(heap)
        
    return nth_min
    
print(nth_smallest([4, 1, 7, 3, 8, 5], 3))   # 4

모듈내 nsmallest() 사용

nsmallest(가져올 수, 적용범위)

from heapq import nsmallest

print(nsmallest(3, [4, 1, 7, 3, 8, 5]))       # [1, 3, 4]
print(nsmallest(3, [4, 1, 7, 3, 8, 5])[-1])   # 4

위의 값이랑은 다르지만 자체모듈을 이용해서 최소값을 가져오는 방법이 있다. 하지만 인덱스[-1] 접근으로 동일한 값을 가져올 수 있다.

모듈내 nlargest() 사용

from heapq import nlargest

print(nlargest(1, [4, 1, 7, 3, 8, 5]))     # [8]
print(nlargest(1, [4, 1, 7, 3, 8, 5])[0])  # 8

마찬가지도 동일한 방법으로 nlargest()를 사용하여 최대값을 구할수 있다.

[응용] heap 정렬

from heapq import heappush, heappop

def heap_sort(nums):
    heap = []
    for num in nums:
        heappush(heap, num)

    sorted_nums = []
    while heap:
        sorted_nums.append(heappop(heap))
    return sorted_nums

print(heap_sort([4, 1, 7, 3, 8, 5]))  # [1, 3, 4, 5, 7, 8]

heappush와 heappop을 통해 정렬을 구현할 수가 있다.

정답

import sys
from heapq import heappush, heappop
n = int(input())
heap = []
for _ in range(n):
    num = int(sys.stdin.readline())
    if num == 0:
        if heap:
            print(heappop(heap)[1])
        else:
            print(0)
    else:
        heappush(heap, (-num, num))