Pythonで二分木の最長連続パスの長さを求めるアルゴリズムと実装

二分木（バイナリツリー）が与えられたとき、木の中にある最長の連続パスの長さを求めることを考えます。ここでの「連続パス」とは、隣り合うノードの値が1ずつ増加、または1ずつ減少していくようなノードの並びのことです。

問題の例

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

この場合、最も長い連続シーケンスは [2, 3, 4, 5, 6] となるため、出力は 5 になります。

解き方のアプローチ

この問題は、再帰的に各ノードを訪問しながら「増加パス」と「減少パス」の長さを追跡することで解けます。手順は以下の通りです。

• ルートがnullの場合は0を返す
• 最大パス長を記録する変数 maxPath を0で初期化する
• ヘルパー関数 helper() を定義する。この関数はノードを受け取り、以下の処理を行う
  • inc（増加パスの長さ）と dec（減少パスの長さ）をそれぞれ1で初期化する
  • 左の子が存在する場合は再帰的に helper() を呼び出し、left_inc と left_dec を取得する。存在しない場合は [0, 0] を設定する
  • 右の子についても同様に right_inc と right_dec を取得する
  • 左の子が存在し、「現在のノードの値 − 左の子の値」が1の場合、inc を inc と left_inc + 1 の大きい方に更新する
  • 同様に差が−1の場合は、dec を dec と left_dec + 1 の大きい方に更新する
  • 右の子に対しても同じ判定を行い、inc / dec を更新する
  • 左右両方の子が存在し、左の子の値が現在のノードより1大きく、かつ現在のノードの値が右の子より1大きい場合（折り返しパス）、maxPath を left_dec + right_inc + 1 と比較して更新する
  • 逆に、左の子の値が現在のノードより1小さく、右の子の値がさらに1小さい場合も、maxPath を left_inc + right_dec + 1 と比較して更新する
  • 最後に maxPath を inc・dec とも比較して最大値を維持し、inc と dec を返す
• メイン処理では helper(root) を呼び出し、maxPath を返す

実装例

それでは、実際のPythonコードを見てみましょう。

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

def print_tree(root):
    if root is not None:
        print_tree(root.left)
        print(root.val, end = ', ')
        print_tree(root.right)

class Solution:
    def solve(self, root):
        if not root:
            return 0
        self.maxPath = 0

        def helper(node):
            inc, dec = 1, 1
            if node.left:
                left_inc, left_dec = helper(node.left)
            else:
                left_inc, left_dec = 0, 0
            if node.right:
                right_inc, right_dec = helper(node.right)
            else:
                right_inc, right_dec = 0, 0

            if node.left and node.val - node.left.val == 1:
                inc = max(inc, left_inc + 1)
            elif node.left and node.val - node.left.val == -1:
                dec = max(dec, left_dec + 1)

            if node.right and node.val - node.right.val == 1:
                inc = max(inc, right_inc + 1)
            elif node.right and node.val - node.right.val == -1:
                dec = max(dec, right_dec + 1)

            if (node.left and node.right and node.left.val - node.val == 1 and node.val - node.right.val == 1):
                self.maxPath = max(self.maxPath, left_dec + right_inc + 1)
            elif (node.left and node.right and node.left.val - node.val == -1
                  and node.val - node.right.val == -1):
                self.maxPath = max(self.maxPath, left_inc + right_dec + 1)

            self.maxPath = max(self.maxPath, inc, dec)
            return inc, dec

        helper(root)
        return self.maxPath

ob = Solution()
root = TreeNode(3)
root.left = TreeNode(2)
root.right = TreeNode(4)
root.right.left = TreeNode(5)
root.right.right = TreeNode(9)
root.right.left.left = TreeNode(6)
print(ob.solve(root))

入力

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

出力

5

計算量について

このアルゴリズムは各ノードを一度だけ訪問するため、時間計算量は O(n)（nはノード数）、再帰呼び出しによる空間計算量は木の高さに依存し、最悪の場合（偏った木）で O(n)、バランスの取れた木では O(log n) となります。

まとめ

本記事では、二分木における最長連続パスの長さを求める方法を解説しました。ポイントは、各ノードごとに「上昇方向のパス長（inc）」と「下降方向のパス長（dec）」を再帰的に計算し、左右の子をつなぐ折り返しパスも考慮して最大値を更新していくことです。この手法を使えば、線形時間で効率よく答えを求められます。