Pythonで合計が偶数となる最長パスの長さを求めるプログラム

問題の概要

二分木（バイナリツリー）が与えられたとき、ノード値の合計が偶数になる最長パスの長さを求めることを考えます。

例えば、次のような木構造が入力として与えられた場合、出力は 5 になります。これはパス [5, 2, 4, 8, 5] を通ったときの合計が 24（偶数）となるためです。

解法のアプローチ：DFS（深さ優先探索）

この問題は、DFSを用いて各ノードから下方向に伸びる「偶数和のパス」と「奇数和のパス」の長さを同時に追跡することで解けます。dfs() 関数はペア (left_0, left_1) を返し、それぞれ「そのノードから下に伸びる合計が偶数のパスの最大長」「合計が奇数のパスの最大長」を表します。

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

• 関数 dfs() を定義します。引数としてノードを受け取ります。
• ノードが null の場合、ペア (0, -inf) を返します。
• (left_0, left_1) := dfs(ノードの左の子)、(right_0, right_1) := dfs(ノードの右の子) を取得します。
• ノードの値が奇数の場合：
  • ans を ans、(left_1 + right_0 + 1)、(left_0 + right_1 + 1) の最大値で更新します。
  • ペア (max(left_1 + 1, right_1 + 1, 0), max(left_0 + 1, right_0 + 1)) を返します。
• ノードの値が偶数の場合：
  • ans を ans、(left_0 + right_0 + 1)、(left_1 + right_1 + 1) の最大値で更新します。
  • ペア (max(left_0 + 1, right_0 + 1, 0), max(left_1 + 1, right_1 + 1)) を返します。
• メインメソッドでは、ans := 0 と初期化し、dfs(root) を呼び出した後、ans を返します。

Pythonでの実装例

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

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 0, float("-inf")
            left_0, left_1 = dfs(node.left)
            right_0, right_1 = dfs(node.right)
            if node.val & 1:
                self.ans = max(self.ans, left_1 + right_0 + 1, left_0 + right_1 + 1)
                return max(left_1 + 1, right_1 + 1, 0), max(left_0 + 1, right_0 + 1)
            else:
                self.ans = max(self.ans, left_0 + right_0 + 1, left_1 + right_1 + 1)
                return max(left_0 + 1, right_0 + 1, 0), max(left_1 + 1, right_1 + 1)

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

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

入力

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

出力

5

まとめ

このアルゴリズムでは、各ノードを一度だけ訪問するため、時間計算量は O(N)、再帰によるスタック使用のため空間計算量も O(N) となります。ノードの値の偶奇に応じて、左右の部分木からのパスを組み合わせたときの合計の偶奇がどう変化するかを丁寧に追跡することがポイントです。