[파이썬] 자료구조 heap 직접 구현해보기

import sys
input = sys.stdin.readline

# idx에 위치하고 있는 새로 추가한 원소를 규칙에 맞는 위치로 이동하기
def up_heapify(idx, queue):
    child_idx = idx
    while child_idx != 0:
        parent_idx = (child_idx - 1) // 2
        if queue[parent_idx] > queue[child_idx]:
            queue[parent_idx], queue[child_idx] = queue[child_idx], queue[parent_idx]
            child_idx = parent_idx
        else:
            return

def heap_push(item, queue):
    queue.append(item)
    up_heapify(len(queue)-1, queue)

def find_smaller_child_idx(idx, queue):
    parent = idx
    size = len(queue)
    left_child = idx*2 + 1
    right_child = idx*2 + 2
    
    if left_child < size and queue[left_child] < queue[parent]:
        parent = left_child
    if right_child < size and queue[right_child] < queue[parent]:
        parent = right_child
    return parent

def down_heapify(idx, queue):
    child_idx = find_smaller_child_idx(idx, queue)
    
    while child_idx != idx:
        queue[child_idx], queue[idx] = queue[idx], queue[child_idx]
        idx = child_idx
        child_idx = find_smaller_child_idx(idx, queue)

def heap_pop(queue):
    if len(queue) == 0:
        return 0
    
    tmp = queue[0]
    queue[0] = queue[-1]
    queue.pop()
    down_heapify(0, queue)
    return tmp

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

for _ in range(N):
    query = tuple(input().split())
    
    if query[0] == "pop":
        print(heap_pop(pq))
    else:
        item = int(query[1])
        heap_push(item, pq)

result

1000
push 4
push 5
push 2
push 3
push 1
pop
1
pop
2
pop
3
pop
4
pop
5

최대힙은 조건문 부등호 등만 바꿔주면 됨