Pythonで二分木の最長交互パス(ジグザグパス)の長さを求めるプログラム

問題概要

二分木が与えられたとき、「左の子 → 右の子 → 左の子…」のように左右交互にたどりながら下へ進む最長のパス(交互パス)の長さを求めます。

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

この場合、交互パスは [2, 4, 5, 7, 8] となるため、出力は 5 になります。

解き方のステップ

この問題を解くには、以下の手順に従います。

• ルートが null(空)の場合は 0 を返します。
• dfs() 関数を定義します。この関数は node(現在のノード)、count(現在のパス長)、flag(次に進むべき方向)を引数に取ります。
• node が null でない場合:
  • flag が True の場合:
    • a ← dfs(node の左の子, count + 1, False)
    • b ← dfs(node の右の子, 1, True)
  • flag が False の場合:
    • a ← dfs(node の右の子, count + 1, True)
    • b ← dfs(node の左の子, 1, False)
  • a と b のうち大きい方を返します。
• node が null の場合は count を返します。
• メインメソッドでは以下を行います。
  • a ← dfs(root の左の子, 1, False)
  • b ← dfs(root の右の子, 1, True)
  • a と b のうち大きい方を返します。

アルゴリズムのポイント

このアプローチでは深さ優先探索(DFS)を活用しています。flag は「次にどちらの子へ移動すれば交互パスを継続できるか」を表しています。

• flag が True のとき: 左の子へ進めば交互パスを継続できるため count + 1 を渡し、右の子へは新たなパスの始点として count を 1 にリセットして渡します。
• flag が False のとき: 右の子へ進めば交互パスを継続できるため count + 1 を渡し、左の子へは count を 1 にリセットして渡します。

各ノードで「交互パスを継続するケース」と「そこから新しいパスを始めるケース」の両方を試し、その最大値を返すことで、木全体の中で最も長い交互パスの長さを求めることができます。

Pythonでの実装例

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

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

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

        def dfs(node, count, flag):
            if node:
                if flag == True:
                    a = dfs(node.left, count + 1, False)
                    b = dfs(node.right, 1, True)
                elif flag == False:
                    a = dfs(node.right, count + 1, True)
                    b = dfs(node.left, 1, False)

                return max(a, b)
            return count

        a = dfs(root.left, 1, False)
        b = dfs(root.right, 1, True)

        return max(a, b)

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

入力

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

出力

5

この例では、パス [2, 4, 5, 7, 8] が最も長いため、答えは 5 となります。