[WEEK 02] 알고리즘 - 우선순위 큐 문제풀이

11279. 최대 힙

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

먼저 내가 heapq 라이브러리를 사용하지 않고 작성한 코드는 당연히 시간 초과가 나왔다.

import sys

input = sys.stdin.readline

N = int(input())

lst = [int(input().split('\n')[0]) for _ in range(N)]
arr = []

for x in lst:
    if x == 0:
        if len(arr) == 0:
            print(0)
        else:
            print(max(arr))
            arr.remove(max(arr))
    else:
        arr.append(x)

다음은 heapq 라이브러리를 이용한 답안 코드이다.

import sys
import heapq

sys.stdin = open("5-1_11279.txt", "r")
input = sys.stdin.readline

N = int(input())
heap = []

for _ in range(N):
    num = int(input())
    if not num:
        # 힙이 비어 있지 않으면 최대값 출력
        if heap:
            print(-(heapq.heappop(heap)))
        # 힙이 비어 있으면 0 출력
        else:
            print(0)
    else:
        heapq.heappush(heap, -num)

heapq.heappush(heap, item)
Push the value item onto the heap, maintaining the heap invariant.

heapq.heappop(heap)
Pop and return the smallest item from the heap, maintaining the heap invariant. If the heap is empty, IndexError is raised. To access the smallest item without popping it, use heap[0].

heapq 라이브러리의 heappush()는 최대 힙 대신 최소 힙을 지원한다. 그렇기 때문에 값을 넣을 때 -를 곱하고 (heappush(heap, -cmd), 출력할 때도 -를 곱하면 (-(heapq.heappop(heap)) 최대 힙을 구할 수 있다.

---

1655. 가운데를 말해요

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

문제에서 주어진 '수의 개수가 짝수개라면 중간에 있는 두 수 중에서 작은 수를 말해야 한다'는 조건은 이 문제의 답을 구하는 데에 핵심적인 조건이다.

작은 입력값들을 저장하는 최대 힙 left_heap과 큰 입력값들을 저장하는 최소 힙 right_heap을 생성하여 left_heap의 0번째 원소(루트 노드이자 최대값)를 출력한다.

import sys
import heapq

input = sys.stdin.readline

N = int(input())

# left_heap: 최대 힙 / right_heap: 최소 힙
left_heap = []
right_heap = []

# 수빈이가 외친 값 입력
for i in range(N):
    num = int(input())
    # left_heap과 right_heap에 번갈아 가며 입력값을 추가
    if i % 2 == 0:
        # 작은 값들이 저장된 left_heap에서 필요한 값은 최댓값이므로 최대 힙
        heapq.heappush(left_heap, -num)
    else:
        # 큰 값들이 저장된 right_heap에서 필요한 값은 최솟값이므로 최소 힙
        heapq.heappush(right_heap, num)
    
    # 만약 left_heap에서 right_heap보다 큰 값이 저장되었으면 바꿔준다.
    if right_heap and -left_heap[0] > right_heap[0]:
        heapq.heappush(left_heap, -heapq.heappop(right_heap))
        heapq.heappush(right_heap, -heapq.heappop(left_heap))
    
    print(-left_heap[0])

11279번 문제에서와 같이 최대 힙이므로 -값으로 입출력한다.

---

1715. 카드 정렬하기

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

카드 묶음의 총 비교 횟수를 최소화하려면 작은 카드 묶음부터 큰 카드 묶음 순으로 비교해야 한다. 작은 카드 묶음 순으로 정렬하기 위해서 원소를 오름차순으로 정렬하는 최소 힙을 이용한다.

import sys
import heapq

input = sys.stdin.readline

N = int(input())
heap = []
heapq.heapify(heap)

for i in range(N):
    heapq.heappush(heap, int(input()))

answer = 0

while len(heap) > 1:
    num1 = heapq.heappop(heap)
    num2 = heapq.heappop(heap)
    sum = num1 + num2
    answer += sum
    heapq.heappush(heap, sum)
    
print(answer)

최대값 또는 최소값을 구하는 문제는 힙을 이용하는 것이 좋다.