ヒープの実装

Complete Binary Tree

n = int(input())

A = [0] + list(map(int,input().split()))

def left(i):
    return A[2*i]

def right(i):
    return A[2*i+1]

def parent(i):
    return A[i//2]

def print_node(i):
    print("node %d: key = %d, " % (i,A[i]), end = "")
    if i//2 >= 1:
        print("parent key = %d, " %parent(i), end = "")
    if 2*i < len(A):
        print("left key = %d, " %left(i), end = "")
    if 2*i + 1 < len(A):
        print("right key = %d, " %right(i), end = "")

for i in range(1, len(A)):
    print_node(i)
    print()

Maximum Heap

H = int(input())
A = [None] + list(map(int,input().split()))

def left(i):
    return 2*i

def right(i):
    return 2*i+1

def parent(i):
    return i//2

def maxHeapify(A,i): #節点iを根とする部分木(i以下の木)がマックスヒープになるように設計する。大きいやつをただ上に持ってきているだけ。
    l = left(i)
    r = right(i)
    #左の子、自分、右の子で値が最大のノードを選ぶ。
    if l <= H and A[l]>A[i]:
        largest = l
    else:
        largest = i
    if r <= H and A[r]>A[largest]:
        largest = r
    if largest != i: #iの子の方が値が大きい場合交換
        A[i],A[largest] = A[largest],A[i]
        maxHeapify(A,largest) #再帰呼び出し

def buildMaxHeap(A): #H//2がラストのノードの親。
    for i in range(H//2,0,-1): #後ろからやっていかないとダメ。頭からやっていたら最後の方に大きいやつがある時に、上に来れない。
        maxHeapify(A,i)
buildMaxHeap(A)
print("",*A[1:])

Priority Queue

def left(i):
    return 2*i

def right(i):
    return 2*i+1

def parent(i):
    return i//2

A = [0]*2000000
H = 0

def maxheapify(S,i): #これの計算量は木の高さに依存するのでO(logH)
     #これはノードiより下のノードはmaxheap条件を満たしているが、ノードiの子供たちとの条件を満たしていない時に使う。
    l = left(i)
    r = right(i)
    global H
    if l <= H and S[i] < S[l]:
        largest = l
    else:
        largest = i
    if r <= H and S[largest] < S[r]:
        largest = r
    if largest != i:
        S[largest],S[i] = S[i],S[largest]
        maxheapify(S,largest)

    
def heapIncreaseKey(A,i,key): #buildする必要ない。親との大小関係を確認するだけでいい。
    if key < A[i]: #必要ぽい。#新しい値<古い値だったら これはそういう仮定
        return
    A[i] = key
    #max heap条件を成立させる。
    while i > 1 and A[parent(i)] < A[i]: #ヒープにおいて左の子と右の子の大小は気にしない。
        A[i],A[parent(i)] = A[parent(i)],A[i]
        i = parent(i)
    
def heapInsert(k):
    global H,A
    H += 1
    A[H] = -1 #新しいキーの値を-1にする。もしA[H]に既に値が入っていたら(ヒープの操作ではextractの時に削除したりしない。)
    heapIncreaseKey(A,H,k)
    
def heapExtractMax(A):
    global H
    if H < 1: #何も入っていないのに入れようとしちゃだめ
        return
    max = A[1]
    A[1] = A[H] #最後のを持ってきて、
    H -= 1
    maxheapify(A,1) #これで十分。root以外はmaxheap条件を満たしているから。
    return max

while True:
    a = input().split()
    if a[0] == "insert":
        heapInsert(int(a[1])) #文字列だと不等号が動かない。
    elif a[0] == "extract":
        print(heapExtractMax(A)) #これだとO(H)だけかかるので全然早くない。
    else:
        break