AOJ ALDS1_9_A Complete Binary Tree

備忘録

問題

http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=ALDS1_9_A&lang=ja

回答

import sys, os, math, bisect, itertools, collections, heapq, queue, copy, array

# from scipy.sparse.csgraph import csgraph_from_dense, floyd_warshall
# from decimal import Decimal
# from collections import defaultdict, deque

sys.setrecursionlimit(10000000)

ii = lambda: int(sys.stdin.buffer.readline().rstrip())
il = lambda: list(map(int, sys.stdin.buffer.readline().split()))
fl = lambda: list(map(float, sys.stdin.buffer.readline().split()))
iln = lambda n: [int(sys.stdin.buffer.readline().rstrip()) for _ in range(n)]

iss = lambda: sys.stdin.buffer.readline().decode().rstrip()
sl = lambda: list(map(str, sys.stdin.buffer.readline().decode().split()))
isn = lambda n: [sys.stdin.buffer.readline().decode().rstrip() for _ in range(n)]

lcm = lambda x, y: (x * y) // math.gcd(x, y)

MOD = 10 ** 9 + 7
MAX = float('inf')

parent = lambda i: math.floor(i / 2)
left = lambda i: 2 * i
right = lambda i: 2 * i + 1


def main():
    if os.getenv("LOCAL"):
        sys.stdin = open("input.txt", "r")

    H = ii()
    heap = il()

    for i in range(H):
        string = 'node {0}: key = {1}, '.format(i + 1, heap[i])
        i += 1
        if parent(i) >= 1: string += 'parent key = {0}, '.format(heap[parent(i) - 1])
        if left(i) <= H: string += 'left key = {0}, '.format(heap[left(i) - 1])
        if right(i) <= H: string += 'right key = {0}, '.format(heap[right(i) - 1])
        print(string)


if __name__ == '__main__':
    main()

考え方

螺旋本より、ヒープ入門問題。
与えられた二分ヒープに対して、
先頭から添え字と親または子の情報を出力する。

基本的には指示通り、添え字iに対して親の添え字を求める床関数(math.floor(i / 2))と、
子の添え字を求める計算(left: 2 * i, right: 2 * i +1)で添え字が存在するか否かを求めればよい。
気を付けなければならないのは、配列のindexと添え字の関連について。
rangeでループを回す場合には、0始まりである点と、
0-basedの配列で二分ヒープを受け取った場合には添え字 != indexである点に注意。