Pythonで二分木のルートからリーフまでの経路数値を合計する方法

0〜9の数字のみを含む二分木があるとします。このとき、ルートからリーフ（葉）までの各経路は、一つの数値を表すことができます。

例えば、次のような二分木を考えてみましょう。

この木には「2 → 1」と「2 → 3」の2つの経路があり、それぞれ 21 と 23 を表します。したがって、出力は 21 + 23 = 44 となります。

解法のアプローチ

この問題は、深さ優先探索（DFS）を使うことで効率的に解けます。手順は以下の通りです。

• 再帰関数 dfs() を作成します。引数としてノードと num（初期値は0）を受け取ります。
• ノードが存在する場合：
  • num := num * 10 + ノードの値 で現在の経路の数値を更新します。
  • もしノードがリーフ（左右の子がどちらも存在しない）であれば、
    • sum := sum + num で合計に加算し、
    • num := num / 10 として値を戻し、
    • 関数から return します。
  • そうでなければ、右の子と左の子に対してそれぞれ dfs() を再帰的に呼び出します。
  • 呼び出し後、num := num / 10 で値を戻して return します。
• 最初に sum := 0 で初期化します。
• ルートノードを渡して dfs() を呼び出します。
• 最後に sum を返します。

Pythonでの実装例

以下のコードで、実際の実装を確認してみましょう。

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

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 sumNumbers(self, root):
        self.sum = 0
        self.dfs(root)
        return self.sum

    def dfs(self, node, num=0):
        if node:
            num = num * 10 + node.data
            if not node.right and not node.left:
                self.sum += num
                num /= 10
                return
            self.dfs(node.right, num)
            self.dfs(node.left, num)
            num /= 10
            return

ob1 = Solution()
tree = make_tree([2,1,3])
print(ob1.sumNumbers(tree))

入力

[2,1,3]

出力

44

コードのポイント

このアルゴリズムの核心は、DFSで各経路を辿りながら、num = num * 10 + ノードの値という計算で桁を組み立てていく点です。リーフノードに到達した時点で、その経路が表す完全な数値が完成するため、それを合計に加算します。

計算量は全ノードを一度ずつ訪問するため O(N)、必要なメモリは再帰の深さに依存し O(H)（Hは木の高さ）となります。LeetCode の「Sum Root to Leaf Numbers（問題129）」としても知られる定番の二分木問題です。