Pythonで二分木の各ノードの値が子ノードの値の合計と一致するか判定するプログラム

二分木が与えられたとき、葉ノードを除くすべてのノードについて、その値が「左の子ノードの値 + 右の子ノードの値」と一致しているかどうかを判定する必要があります。

例えば、次のような二分木が入力として与えられた場合を考えてみましょう。

この場合、根ノード18 = 8 + 10、ノード8 = 3 + 5 というように、すべての内部ノードで条件が成り立っているため、出力は True になります。

解決のアプローチ

この問題は、DFS（深さ優先探索）を使って木を再帰的に走査することで解決できます。手順は以下の通りです。

• dfs() 関数を定義します。引数として root を受け取ります。
• root が null（None）の場合は True を返します。
• root の左の子も右の子も null の場合（つまり葉ノードの場合）は True を返します。
• left を 0 で初期化し、root の左の子が存在する場合はその値を代入します。
• right を 0 で初期化し、root の右の子が存在する場合はその値を代入します。
• 「left + right が root の値と等しい」かつ「dfs(左の子) が True」かつ「dfs(右の子) が True」である場合に True を返します。
• メインメソッドでは dfs(root) の結果を返します。

それでは、実際の実装を見てみましょう。

実装例

class TreeNode:
    def __init__(self, data, left=None, right=None):
        self.val = data
        self.left = left
        self.right = right

class Solution:
    def solve(self, root):
        def dfs(root):
            if root == None:
                return True
            if root.left == None and root.right == None:
                return True
            left = 0
            if root.left:
                left = root.left.val
            right = 0
            if root.right:
                right = root.right.val
            return (left + right == root.val) and dfs(root.left) and dfs(root.right)
        return dfs(root)

ob = Solution()
root = TreeNode(18)
root.left = TreeNode(8)
root.right = TreeNode(10)
root.left.left = TreeNode(3)
root.left.right = TreeNode(5)
print(ob.solve(root))

入力

root = TreeNode(18)
root.left = TreeNode(8)
root.right = TreeNode(10)
root.left.left = TreeNode(3)
root.left.right = TreeNode(5)

出力

True

コードのポイント

• 葉ノードの扱い： 葉ノードには子が存在しないため、条件判定の対象外として常に True を返します。
• null チェック： 子ノードが片方しか存在しない場合でもエラーにならないよう、存在しない側は 0 として扱います。
• 再帰的な検証： 各ノードで条件を満たしているかを再帰的に確認するため、木全体の一貫性を保証できます。

このアルゴリズムの計算量は、木の全ノードを一度ずつ訪問するため O(n)（n はノード数）です。空間計算量も再帰呼び出しの深さに依存し、最悪ケースで O(n) となります。