Pythonで二分木をジグザグ（左右交互）にレベル順トラバースする方法

二分木が与えられたとき、各レベルの値を「左から右」「右から左」と交互に切り替えながら出力することを考えます。これはいわゆるジグザグ（鋸歯状）レベル順トラバースと呼ばれる問題です。

例えば、次のような二分木が入力だとします。

この場合、出力は [5, -10, 4, -2, -7, 15] となります。

アルゴリズムの考え方

この問題は、スタックを2つ使うことで効率的に解けます。奇数レベルと偶数レベルで子ノードの追加順序を入れ替えるのがポイントです。手順は以下の通りです。

• root が null の場合は、空のリストを返す
• s1 := root を格納したリストを作成
• s2 := 空のリストを作成
• res := 結果を格納する空のリストを作成
• s1 または s2 が空でない限り、以下を繰り返す
  • s1 が空でない間、以下を繰り返す
    • node := s1 の末尾（最後）の要素を取り出す
    • node の左の子が存在すれば、s2 の末尾に追加
    • node の右の子が存在すれば、s2 の末尾に追加
    • node の値を res の末尾に追加
  • s2 が空でない間、以下を繰り返す
    • node := s2 の末尾（最後）の要素を取り出す
    • node の右の子が存在すれば、s1 の末尾に追加
    • node の左の子が存在すれば、s1 の末尾に追加
    • node の値を res の末尾に追加
• res を返す

s1 から取り出すときは「左→右」の順で s2 へ積み、s2 から取り出すときは「右→左」の順で s1 へ積むことで、探索方向が自動的に交互に反転します。

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):
      if not root:
         return []
      s1 = [root]
      s2 = []
      res = []
      while s1 or s2:
         while s1:
            node = s1.pop()
            if node.left:
               s2.append(node.left)
            if node.right:
               s2.append(node.right)
            res.append(node.val)
         while s2:
            node = s2.pop()
            if node.right:
               s1.append(node.right)
            if node.left:
               s1.append(node.left)
            res.append(node.val)
      return res

ob = Solution()
root = TreeNode(5)
root.left = TreeNode(4)
root.right = TreeNode(-10)
root.left.left = TreeNode(-2)
root.right.left = TreeNode(-7)
root.right.right = TreeNode(15)
print(ob.solve(root))

入力

root = TreeNode(5)
root.left = TreeNode(4)
root.right = TreeNode(-10)
root.left.left = TreeNode(-2)
root.right.left = TreeNode(-7)
root.right.right = TreeNode(15)

出力

[5, -10, 4, -2, -7, 15]

計算量について

このアルゴリズムでは、各ノードをちょうど1回ずつ処理するため、時間計算量は O(n)（n はノード数）、空間計算量も O(n) となります。最悪の場合、最下層のノード数がリストに保持されるためです。

スタック2つを使うこの手法は、deque（両端キュー）やフラグによる方向管理でも同様に実装できますが、スタック方式はコードがシンプルで直感的に理解しやすいのが特徴です。