kb84tkhrのブログ

何を書こうか考え中です あ、あと組織とは関係ないってやつです 個人的なやつ

ALDS1_7_B: Binary Tree

今度は二分木
どっちかというと左子右兄弟表現とやらよりこっちのほうがシンプルで
先に出てきそうな気がしますがどうなんでしょう

同じようなことを書くのですいすい書けます

from sys import stdin

class Node():
    NIL = -1

    def __init__(self, p=NIL, l=NIL, r=NIL, d=-1, h=-1):
        self.p = p
        self.l = l
        self.r = r
        self.d = d
        self.h = h

    def __repr__(self):
        return "Node({}, {}, {}, {})".format(self.p, self.l, self.r, self.d)

def root_of(T):
    return [u.p for u in T].index(Node.NIL)

def sibling_of(T, id):
    p = T[T[id].p]
    return p.r if p.l == id else p.l

def degree_of(T, id):
    l = 0 if T[id].l == Node.NIL else 1
    r = 0 if T[id].r == Node.NIL else 1
    return l + r

def calc_depth(T):

    def rec(id, depth):
        nonlocal T

        T[id].d = depth
        if T[id].l != Node.NIL:
            rec(T[id].l, depth + 1)
        if T[id].r != Node.NIL:
            rec(T[id].r, depth + 1)

    rec(root_of(T), 0)

def calc_height(T):

    def rec(id):
        nonlocal T

        l = 0 if T[id].l == Node.NIL else rec(T[id].l) + 1
        r = 0 if T[id].r == Node.NIL else rec(T[id].r) + 1
        T[id].h = max(l, r)
        return T[id].h

    rec(root_of(T))

def type_of(T, id):
    if T[id].p == Node.NIL:
        return "root"
    elif T[id].l == Node.NIL and T[id].r == Node.NIL:
        return "leaf"
    else:
        return "internal node"

def read_nodes(T):
    for line in stdin:
        id, l, r = [int(x) for x in line.split()]
        T[id].l = l
        if l != Node.NIL:
            T[l].p = id
        T[id].r = r
        if r != Node.NIL:
            T[r].p = id

def print_nodes(T):
    for id in range(len(T)):
        print("node {}: ".format(id), end="")
        print("parent = {}, ".format(T[id].p), end="")
        print("sibling = {}, ".format(sibling_of(T, id)), end="")
        print("degree = {}, ".format(degree_of(T, id)), end="")
        print("depth = {}, ".format(T[id].d), end="")
        print("height = {}, ".format(T[id].h), end="")
        print(type_of(T, id))

def main():
    n = int(stdin.readline())
    T = [Node() for _ in range(n)]

    read_nodes(T)
    calc_depth(T)
    calc_height(T)
    print_nodes(T)

main()

そんなふうに考えていた時期が俺にもありました
read_nodesTの最後の要素だけなんか意味のわからない値が入って
だいぶ困りました

何も考えずにT[l].p = idとかT[r].p = idとかやってたせいでした
lrが-1のとき、配列の末尾にアクセスしてることに気づいてませんでした
Out of rangeとかでエラーにならないの?とか思ってしまいました
あーなりませんね

sibling_ofのところ、解説のソースよりも少し単純になってます
これでは不足かなあ?
親の左か右かは確実に自分のはず
ACにはなりました