백준 11286 - 절댓값 힙 (파이썬)

Answer

import sys

def more(a, b):
    if abs(a) > abs(b):
        return b
    elif abs(b) > abs(a):
        return a
    return min(a, b)

def insert(heap, inp):
    heap.append(inp)
    i = len(heap) - 1
    while i > 1:
        if more(heap[i], heap[i // 2]) == heap[i // 2]:
            break
        heap[i], heap[i // 2] = heap[i // 2], heap[i]
        i = i // 2

def remove(heap):
    if len(heap) == 1:
        return 0
    val = heap[1]
    tmp = heap.pop()
    
    parent = 1
    child = 2
    
    while child <= len(heap) -1:
        if child < len(heap) - 1 and more(heap[child], heap[child +1]) == heap[child + 1]:
            child += 1

        if more(heap[child], tmp) == tmp:
            break

        heap[parent] = heap[child]
        
        parent = child
        child *= 2
        
    if parent <= len(heap) - 1:
        heap[parent] = tmp

    return val

n = int(sys.stdin.readline())
heap = [0]

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

    if num == 0:
        print(remove(heap))
    else:
        insert(heap, num)