Pythonで二分木の各ノードを左右の部分木の合計値で更新するプログラム

問題の概要

二分木が与えられたとき、各ノードの値を「自分自身の値 ＋ 左右の部分木の合計」に置き換えた木を求めることを考えます。つまり、木を葉から根へ（ポストオーダー）たどりながら、すべてのノードを部分木の総和で更新していく処理です。

例えば、次のような二分木が入力として与えられたとします。

この場合、出力は次のようになります。

アルゴリズム

この問題は再帰を使ったポストオーダー走査でシンプルに解けます。手順は以下の通りです。

• 関数 tree_sum() を定義します。引数として木のルートを受け取ります。
• ルートが None（空）の場合は 0 を返します。
• そうでなければ、ルートの値を次のように更新します。
  root.data = tree_sum(root.left) + tree_sum(root.right) + root.data
• 更新後の root.data を呼び出し元に返します。
• メイン側では tree_sum(root) を実行した後、更新済みの root を返します。

Pythonでの実装例

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


def inorder(root):
    if root:
        inorder(root.left)
        print(root.data, end=', ')
        inorder(root.right)


class Solution:
    def solve(self, root):
        def tree_sum(root: TreeNode):
            if root is None:
                return 0
            root.data = tree_sum(root.left) + tree_sum(root.right) + root.data
            return root.data

        tree_sum(root)
        return root


ob = Solution()
root = TreeNode(2)
root.left = TreeNode(3)
root.right = TreeNode(4)
root.left.left = TreeNode(9)
root.left.right = TreeNode(7)

ob.solve(root)
inorder(root)

入力

root = TreeNode(2)
root.left = TreeNode(3)
root.right = TreeNode(4)
root.left.left = TreeNode(9)
root.left.right = TreeNode(7)

出力

9, 19, 7, 25, 4,

動作の解説

コードの流れを具体的に見てみましょう。

• 葉ノード 9 と 7 は子を持たないため、値は変わりません。
• ノード 3 は 3 + 9 + 7 = 19 に更新されます。
• ノード 4 も葉なので 4 のままです。
• ルート 2 は 2 + 19 + 4 = 25 に更新されます。

最後に inorder()（中間順走査）で表示すると、「9, 19, 7, 25, 4,」という結果が得られます。

計算量

• 時間計算量： O(n) ― 各ノードをちょうど1回ずつ訪問するためです。
• 空間計算量： O(h) ― 再帰呼び出しによるスタック消費で、h は木の高さです。木が偏って鎖状になっている場合は O(n) になります。