Pythonで二分木のルートから葉までの各パスが表す数値の合計を求める方法

問題概要

各ノードが0〜9までの一桁の数字を持つ二分木を考えます。ルートから葉へと辿る各パスは、通過したノードの数字を順番につなげた一つの数値を表します。この記事では、木の中のすべてのパスが表す数値の合計を求めるPythonプログラムを解説します。

具体例

例として、次のような二分木を入力とした場合を考えます。

この木には次の3つのパスが存在します。

• 46（4 → 6）
• 432（4 → 3 → 2）
• 435（4 → 3 → 5）

これらの合計は 46 + 432 + 435 = 913 となるため、出力は913になります。

解法のアプローチ

この問題は、深さ優先探索（DFS）を再帰的に行うことで解決できます。手順は以下の通りです。

1. solve() 関数を定義します。引数は現在のノード root と、ここまで辿った数字をつなげた文字列 string（デフォルトは空文字列）です。
2. root が存在し、かつ左右どちらの子も持たない（葉ノードである）場合は、string + str(root.val) を整数に変換して返します。これで一つのパスが完成します。
3. 合計値 total を0で初期化します。
4. root.left が存在する場合、total に solve(root.left, string + str(root.val)) の戻り値を加算します。
5. root.right が存在する場合も同様に加算します。
6. 最後に total を返します。

実装例

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, string=""):
        # 葉ノードに到達したら、パス全体の数字列を整数にして返す
        if root and not root.left and not root.right:
            return int(string + str(root.val))

        total = 0
        if root.left:
            total += int(self.solve(root.left, string + str(root.val)))

        if root.right:
            total += int(self.solve(root.right, string + str(root.val)))

        return total

ob = Solution()
root = TreeNode(4)
root.left = TreeNode(6)
root.right = TreeNode(3)
root.right.left = TreeNode(2)
root.right.right = TreeNode(5)
print(ob.solve(root))

入力

root = TreeNode(4)
root.left = TreeNode(6)
root.right = TreeNode(3)
root.right.left = TreeNode(2)
root.right.right = TreeNode(5)

出力

913

まとめ

このように、再帰的なDFSを使いながらパス上の数字列を組み立てていき、葉ノードに到達した時点で整数に変換して合計することで、すべてのパスが表す数の総和を求められます。各ノードを一度だけ訪問すればよいため、非常にシンプルかつ効率的な実装になっています。文字列の連結ではなく「現在値 × 10 + ノードの値」という数値計算で処理すれば、さらにメモリ効率を高めることも可能です。