Pythonで二分木のノードとその子孫の最大絶対差を求めるプログラム

問題概要

二分木が与えられたとき、任意のノードとその子孫との間の絶対差の最大値を求めることを考えます。

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

この場合、ノード8とノード1の間の差が最も大きくなるため、出力は 7 となります。

解法のアプローチ：DFSを使った追跡

この問題は、DFS（深さ優先探索）を用いることで効率的に解けます。各ノードについて「その部分木内の最小値」と「最大値」を追跡しながら、現在のノードの値との差を順次更新していくのがポイントです。

具体的な手順は以下の通りです。

• dfs() 関数を定義します。引数としてノードを受け取ります。
• ノードが null（空）の場合は、正の無限大と負の無限大からなるリスト [inf, -inf] を返します（初期値として機能させます）。
• left := 左の子ノードに対する dfs() の結果
• right := 右の子ノードに対する dfs() の結果
• res := [left[0]、right[0]、ノードの値の中で最小のもの] と [left[1]、right[1]、ノードの値の中で最大のもの] のペア
• ans := ans、(ノードの値 − res[0])、(res[1] − ノードの値) のうち最大の値に更新
• res を返します。

メインメソッド側では、以下の処理を行います。

• ans := 0 で初期化する
• ルートノードに対して dfs() を呼び出す
• ans を返す

実装例

理解を深めるために、実際のコードを見てみましょう。

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):
        def dfs(node):
            if not node:
                return [float("inf"), float("-inf")]
            left = dfs(node.left)
            right = dfs(node.right)
            res = [min(left[0], right[0], node.val),
                   max(left[1], right[1], node.val)]
            self.ans = max(self.ans, node.val - res[0], res[1] - node.val)
            return res

        self.ans = 0
        dfs(root)
        return self.ans

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

入力

root = TreeNode(1)
root.left = TreeNode(5)
root.right = TreeNode(3)
root.right.left = TreeNode(2)
root.right.right = TreeNode(8)
root.right.left.left = TreeNode(7)
root.right.left.right = TreeNode(4)

出力

7

アルゴリズムのポイント

このコードがうまく動作する理由は、葉に到達したときのベースケースとして [float("inf"), float("-inf")] を返す点にあります。これにより、空の部分木が最小値・最大値の計算に影響を与えることなく、各ノードで正しく比較が行えます。

また、res[0] には部分木全体の最小値、res[1] には最大値が格納されるため、「ノードの値 − 最小値」と「最大値 − ノードの値」のどちらか大きい方を ans に反映していくことで、すべてのノードと子孫の組み合わせの中から最大の絶対差を確実に見つけられます。

計算量は、全ノードを一度ずつ訪問するため時間計算量 O(n)、再帰の深さは木の高さに依存するため空間計算量 O(h)（h は木の高さ）となります。非常に効率的な解法と言えるでしょう。