Pythonで二分木のパス合計（Path Sum）を判定する方法

パス合計問題とは

二分木と目標の合計値が与えられたとき、根から葉までの経路をたどった際のノード値の合計が、与えられた値と一致するような経路が存在するかどうかを判定します。

例として、木が [0, -3, 9, -10, null, 5] という構成で、合計値が 14 の場合を考えてみましょう。このとき、0 → 9 → 5 という経路が存在し、その合計はちょうど 14 になるため、答えは True となります。

解法のアプローチ

この問題は再帰を使うことで簡潔に解くことができます。手順は以下の通りです。

• 根ノードが null（空）の場合、False を返します。
• 左右の子ノードが両方とも空（つまり葉ノード）の場合、sum - root.val == 0 であれば True、そうでなければ False を返します。
• それ以外の場合は、左部分木と右部分木に対して残りの合計値（sum - root.val）を渡して再帰的に呼び出し、その結果を OR 演算で結合して返します。

実装例

以下のコードを実行すると、実際の動作を確認できます。

# 二分木のノード定義
class TreeNode(object):
    def __init__(self, x):
        self.data = x
        self.left = None
        self.right = None

def insert(temp, data):
    que = []
    que.append(temp)
    while (len(que)):
        temp = que[0]
        que.pop(0)
        if (not temp.left):
            if data is not None:
                temp.left = TreeNode(data)
            else:
                temp.left = TreeNode(0)
            break
        else:
            que.append(temp.left)
        if (not temp.right):
            if data is not None:
                temp.right = TreeNode(data)
            else:
                temp.right = TreeNode(0)
            break
        else:
            que.append(temp.right)

def make_tree(elements):
    Tree = TreeNode(elements[0])
    for element in elements[1:]:
        insert(Tree, element)
    return Tree

class Solution(object):
    def hasPathSum(self, root, sum):
        """
        :type root: TreeNode
        :type sum: int
        :rtype: bool
        """
        if not root:
            return False
        if not root.left and not root.right and root.data is not None:
            return sum - root.data == 0
        if root.data is not None:
            return self.hasPathSum(root.left, sum-root.data) or self.hasPathSum(root.right, sum-root.data)

tree1 = make_tree([0,-3,9,-10,None,5])
ob1 = Solution()
print(ob1.hasPathSum(tree1, 14))

入力

tree1 = make_tree([0,-3,9,-10,None,5])

出力

True

処理の流れの解説

このアルゴリズムでは、各ノードを訪問するたびに、現在の合計値からそのノードの値を引いていきます。葉ノードに到達した時点で、残りの合計値が 0 になっていれば、目的の経路が存在することになります。計算量は二分木の全ノード数を n とすると O(n)、最悪の場合（木が偏っている場合）の空間計算量は O(n) となります。