Python 題目

PROBLEM

樹的屬性計算

一題一題刷

概念

計算樹中每個節點的重要屬性:

程式碼實作

from collections import defaultdict

def dfs(Tree, now_node, depth):
    depths[now_node] = depth  # 記錄當前節點深度
    height = 0
    
    # 遞迴計算所有子節點的高度
    for next_node in Tree[now_node]:
        height = max(dfs(Tree, next_node, depth + 1), height)
    
    heights[now_node] = height  # 記錄當前節點高度
    return height + 1  # 回傳給父節點

# 初始化
tree = defaultdict(list)
n = int(input())
parent = [-1] * n  # parent[i] = i 的父節點
depths = [0] * n   # depths[i] = i 的深度
heights = [0] * n  # heights[i] = i 的高度

# 讀取輸入
for _ in range(n):
    id, k, *c = map(int, input().split())
    for i in c:
        parent[i] = id  # 設定子節點的父節點
    tree[id].extend(sorted(c))  # 將子節點排序後加入

# 找到根節點並開始 DFS
root = parent.index(-1)
dfs(tree, root, 0)

# 輸出結果
for id in range(n):
    types = ""
    if parent[id] == -1:
        types = "root"
    elif len(tree[id]) > 0:
        types = "internal node"
    else:
        types = "leaf"
    
    print(f"node {id}: parent = {parent[id]}, depth = {depths[id]}, "
          f"height = {heights[id]}, {types}, {tree[id]}")
import sys
sys.setrecursionlimit(1000000)  # 增加遞迴深度限制

當樹的深度很大時,Python 預設的遞迴深度可能不夠,需要手動增加限制。