Pythonで実装する二分木のジグザグレベル順走査（Zigzag Level Order Traversal）

二分木が与えられたとき、そのジグザグレベル順走査（Zigzag Level Order Traversal）の結果を求める問題を考えます。これは、第1レベルは左から右へ、第2レベルは右から左へ、第3レベルは再び左から右へ……というように、階層ごとに走査の向きを交互に切り替えながらノードを訪問する手法です。

例として、次のような二分木を扱います。

この木に対する走査結果は [[3], [20, 9], [15, 7]] になります。ルートの 3 を含む第1レベル、右から左へ読む第2レベル（20, 9）、そして左から右へ読む第3レベル（15, 7）という具合です。

アルゴリズムの流れ

キュー（FIFO）と走査方向を表すフラグを使い、レベルごとに次の手順で処理を進めます。

1. 木が空の場合は、空のリストを返します。
2. キューを作成してルートを追加し、各レベルのノードを保持する res、最終結果となる値のリスト res2、走査方向を示すフラグ flag = True を用意します。
3. キューが空になるまで、以下を繰り返します。
   • 現在キューに入っているノードの一覧を res に、それらの値を res2 に追加します。
   • flag が True の場合（次のレベルを右から左へ読む）：現在のレベルを末尾のノードから順に調べ、右の子 → 左の子の順で新しいキューに追加します。
   • flag が False の場合（次のレベルを左から右へ読む）：同様に末尾のノードから順に、今度は左の子 → 右の子の順で追加します。
   • キューを新しいレベルのノードで置き換え、flag を反転させます。
4. すべてのレベルを処理し終えたら、res2 を返します。

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 = [temp]
    while que:
        temp = que.pop(0)
        if not temp.left:
            temp.left = TreeNode(data if data is not None else 0)
            break
        que.append(temp.left)
        if not temp.right:
            temp.right = TreeNode(data if data is not None else 0)
            break
        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 zigzagLevelOrder(self, root):
        if not root:
            return []
        queue = [root]
        res = []       # 各レベルのノードオブジェクトを保持
        res2 = []      # 最終的に出力する値のリスト
        flag = True    # True: 次のレベルは右から左へ
        while queue:
            res.append(list(queue))
            res2.append([node.data for node in queue if node.data != 0])
            next_queue = []
            if flag:
                # 末尾のノードから順に、右の子→左の子の順で追加
                for node in reversed(queue):
                    if node.right:
                        next_queue.append(node.right)
                    if node.left:
                        next_queue.append(node.left)
            else:
                # 末尾のノードから順に、左の子→右の子の順で追加
                for node in reversed(queue):
                    if node.left:
                        next_queue.append(node.left)
                    if node.right:
                        next_queue.append(node.right)
            queue = next_queue
            flag = not flag
        return res2


ob = Solution()
tree = make_tree([3, 9, 20, None, None, 15, 7])
print(ob.zigzagLevelOrder(tree))

なお、この実装では存在しないノードを値 0 のダミーノードとして表現しているため、結果を組み立てる際に data != 0 の条件でダミーを除外しています。

入力

[3,9,20,null,null,15,7]

出力

[[3], [20, 9], [15, 7]]

計算量

各ノードをちょうど一度ずつ処理するため、時間計算量は O(n) です。また、各レベルのノードをキューとリストに保持する必要があるため、空間計算量も O(n) となります（n はノードの総数）。