Pythonで二分木の各レベルの葉ノードの合計値の積を求める方法

問題の概要

二分木（バイナリツリー）が与えられたとき、以下の操作を実行することを考えます。

• 各レベルについて、そのレベルに葉ノードが存在する場合はすべての葉ノードの値の合計を求めます。葉ノードが存在しないレベルは無視します。
• 求めたすべての合計値の積を計算して返します。

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

この場合、出力は 270 になります。最初の2つのレベルには葉ノードが存在しません。第3レベルには葉ノードが1つだけあり、その値は 9 です。最後のレベルには 2、12、5、11 という4つの葉ノードがあります。したがって、結果は 9 × (2 + 12 + 5 + 11) = 270 となります。

解法のアルゴリズム

この問題を解くためには、幅優先探索（BFS）の考え方を用いて、以下の手順で処理を進めます。

1. ルートが null の場合は 0 を返します。
2. 結果を格納する変数 res を 1 で初期化します。
3. キューを作成し、ルートを末尾に追加します。
4. 以下の処理を無限に繰り返します。
   • キューのサイズを no_of_nodes とします。
   • no_of_nodes が 0 の場合はループを抜けます。
   • sum_level を 0、found_leaf を False で初期化します。
   • no_of_nodes が 0 より大きい間、以下を繰り返します。
     • キューの先頭要素を curr_node として取得します。
     • curr_node が葉ノードの場合は、found_leaf を True にし、sum_level に curr_node の値を加算します。
     • キューの先頭要素を削除します。
     • curr_node の左の子が存在すればキューに追加します。
     • curr_node の右の子が存在すればキューに追加します。
     • no_of_nodes を 1 減らします。
   • found_leaf が True の場合、res に sum_level を掛けます。
5. res を返します。

Pythonでの実装例

理解を深めるために、以下の実装例を見てみましょう。

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

def isLeaf(root):
    return (not root.left and not root.right)

def find_res(root):
    if (not root):
        return 0
    res = 1
    que = []
    que.append(root)
    while (True):
        no_of_nodes = len(que)
        if (no_of_nodes == 0):
            break
        sum_level = 0
        found_leaf = False
        while (no_of_nodes > 0):
            curr_node = que[0]
            if (isLeaf(curr_node)):
                found_leaf = True
                sum_level += curr_node.data
            que.pop(0)
            if (curr_node.left != None):
                que.append(curr_node.left)
            if (curr_node.right != None):
                que.append(curr_node.right)
            no_of_nodes -= 1
        if (found_leaf):
            res *= sum_level
    return res

root = TreeNode(8)
root.left = TreeNode(8)
root.right = TreeNode(6)
root.left.right = TreeNode(7)
root.left.left = TreeNode(9)
root.left.right.left = TreeNode(2)
root.left.right.right = TreeNode(12)
root.right.right = TreeNode(10)
root.right.right.left = TreeNode(5)
root.right.right.right = TreeNode(11)
print(find_res(root))

入力

root = TreeNode(8)
root.left = TreeNode(8)
root.right = TreeNode(6)
root.left.right = TreeNode(7)
root.left.left = TreeNode(9)
root.left.right.left = TreeNode(2)
root.left.right.right = TreeNode(12)
root.right.right = TreeNode(10)
root.right.right.left = TreeNode(5)
root.right.right.right = TreeNode(11)

出力

270

まとめ

このアルゴリズムでは、キューを利用したレベル順走査（BFS）によって各レベルのノードを効率的に処理しています。各レベルで葉ノードが見つかった場合のみ合計値を計算し、最終的にすべての合計値の積を返すことで、目的の結果を得ることができます。計算量は木のノード数 n に対して O(n) となり、空間計算量も O(n) であるため、大規模な二分木にも対応できる効率的な解法です。