Pythonで解く二分木の最大パス和（Maximum Path Sum）

問題の概要

空でない二分木が1つ与えられます。この木における「最大パス和」を求めるのが目的です。ここでいうパスとは、あるノードを起点として、親子関係で結ばれたノードをたどり任意のノードへ至るまでのノード列のことです。パスには少なくとも1つのノードが含まれている必要がありますが、必ずしも根（ルート）ノードを通る必要はありません。

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

この場合の出力は 32 となります。

アルゴリズムの考え方

各ノードを「パスの折り返し地点」として捉えるのがポイントです。あるノードを頂点とするパスの和は、「左部分木からの最大寄与 ＋ 右部分木からの最大寄与 ＋ そのノード自身の値」で表せます。これを再帰的に計算しながら、全体の最大値を更新していきます。負の値を持つ部分木はパスの和を下げるだけなので、max(0, ...) によって切り捨てるのが工夫点です。

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

• solve(node) というメソッドを定義する
• node が null、または node の値が 0 の場合は 0 を返す
• left := max(0, solve(node の左の子)) を計算する
• right := max(0, solve(node の右の子)) を計算する
• ans := max(ans, left + right + node の値) で答えを更新する
• node の値 + max(left, right) を返す（親へ渡せるのは左右どちらか片方のみのため）
• メイン側では ans を -∞ で初期化し、solve(root) を呼び出した後、ans を返す

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 maxPathSum(self, root):
        self.ans = -float('inf')
        self.solve(root)
        return self.ans

    def solve(self, node):
        if not node or node.data == 0:
            return 0
        left = max(0, self.solve(node.left))
        right = max(0, self.solve(node.right))
        self.ans = max(self.ans, left + right + node.data)
        return node.data + max(left, right)

ob = Solution()
root = make_tree([-10,9,10,None,None,15,7])
print(ob.maxPathSum(root))

入力

[-10,9,10,None,None,15,7]

出力

32

計算量について

このアルゴリズムは各ノードを一度だけ訪問するため、時間計算量は O(n)（n はノード数）です。空間計算量は再帰呼び出しの深さに依存し、バランスの取れた木では O(log n)、片寄った木では最悪 O(n) となります。負の寄与を切り捨てる処理により、すべてのノードを起点としたパスの中から効率よく最大値を見つけられる点が、この手法の優れたところです。