TIL 코딩공장

백준 파이썬 공부 2025.03.05
11279번 최대 힙 (실버 2)

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

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

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

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

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

4. 문제 풀이
처음으로 그냥 리스트 정렬을 이용해서 풀어보겠다.

>>> 코드 0. 리스트 정렬 >> 시간초과

import sys

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

당연하지만 시간초과가 된다. 정렬을 매번 시행하면 시간이 너무 많이 걸린다.

하나는 파이썬에서 제공하는 힙 라이브러리를 이용하여 쉽게 짜보았다.
파이썬 heapq 라이브러리에서 제공하는 힙을 최소힙이기 때문에,
최대 힙으로 만들려면 약간의 꼼수가 필요하다.
해당 값을 저장 할 때, (-n, n)으로 해당 값의 음수 버전을 같이 튜플로 묶어 힙에 저장하면,
최대 힙을 구현 할 수 있다.

>>>코드 1. heapq 라이브러리 사용

import sys
import heapq

def heap(n):
    if n == 0:
        if len(num) == 0:
            print(0)
        else:
            a, b = heapq.heappop(num)
            print(b)
    else:
        heapq.heappush(num, (-n, n))

num = []
for _ in range(int(sys.stdin.readline())):
    n = int(sys.stdin.readline())
    heap(n)

다음은 그냥 쌩으로 힙을 구현해보겠다.

<<힙>>
- 우선순위 큐를 위해 고안된 완전 이진 트리 형태의 자료구조
- 여러 값 중, 최대 값과 최소값을 빠르게 찾아내도록 만들어진 자료 구조
- 반 정렬 상태를 유지
- 이진 탐색 트리와 달리 중복된 값이 허용됨
- 최대 힙과 최소 힙이 존재

최대 힙 : 맨 위에 가장 큰 값 존재, 인덱스 1에 최대값 존재
최소 힙 : 맨 위에 가장 작은 값 존재, 인덱스 1에 최소값 존재

저장 방식
- 새 노드의 우선순위를 가장 낮게 생각하여 맨 끝에 위치
- 부모 노드와 우선순위를 비교하여 위치 변경
- 계속 반복

삭제 방식
- 루트 노드(우선순위가 가장 높은 노드)을 제거
- 우선순위 가장 낮은 노드를 루트 노드 자리로 이동 
- 2개의 자식 노드들 중 값이 더 큰 것과 비교하여 위치 변경
- 계속 반복

알아야 할 것은
우선 순위가 가장 높은 노드(루트노드)를 뺴면, 
딱히 오름차순이나 내림차순으로 정렬되어있지 않다는 것이다. 
그냥 규칙에 따라 저장하고 삭제하고 있을 뿐, 나머지의 순서가 크기순은 아니다.

>>> 코드 2. 최대 힙 구현

import sys

def pop_heap():
    if len(num) == 1:
        return num.pop()
    
    pop_data, num[0] = num[0], num.pop()
    
    if num:
        current, child = 0, 1
        length = len(num)
        
        while child < length:
            sibling = child + 1
        
            if sibling < length and num[child] < num[sibling]:
                child = sibling
            if num[current] < num[child]:
                num[current], num[child] = num[child], num[current]
                current = child
                child = current * 2 + 1
            else:
                break
    
    return pop_data
    
def push_heap(n):
    num.append(n)
    current = len(num)-1 
    parent = (current-1) // 2
    
    while current > 0:
        if num[current] > num[parent]:
            num[current], num[parent] = num[parent], num[current]
            current = parent
            parent = (current-1) // 2  
        else:
            break

def heap(n):
    if n == 0:
        if len(num) == 0:
            print(0)
        else:
            print(pop_heap())
    else:
        push_heap(n)
    
num = []
for _ in range(int(sys.stdin.readline())):
    n = int(sys.stdin.readline())
    heap(n)

5. 문제 링크
https://www.acmicpc.net/problem/11279