Swift教室 Swift と競技プログラミングの教室

26. Binary Tree Level Order Traversal

読了目安 約1分

Tree / LC 102 (Medium) — レベルごとの走査。BFS の基本形。

この章の目次

問題

LeetCode 102. Binary Tree Level Order Traversal(Medium / カテゴリ: Tree)

二分木の根 root が与えられます。 同じ深さのノードの値を左から右に並べた配列を、浅いレベルから順に集めた 2 次元配列を返します。 例えば根が 3、その子が 9 と 20 なら、答えの先頭は [[3], [9, 20], ...] です。

考え方

ヒント 1

深さ方向に潜る再帰ではなく、同じ深さのノードをまとめて処理する幅優先探索(BFS)が素直です。

ヒント 2

「現在のレベルのノードの配列」を持ちます。 値を記録したら、全ノードの子を左から順に集めて次のレベルの配列を作ります。 配列が空になるまで繰り返します。

Swift 実装のポイント

  • Swift には標準のキューが無いため、「今のレベルの配列 → 次のレベルの配列」を作り直す書き方が簡単です。先頭削除 removeFirst()(O(N))を避けられます。
  • 値の取り出しは level.map { $0.val } で 1 行になります。
  • 空の木は guard let で最初に処理して [] を返します。
模範解答
Swift
class Solution {
    func levelOrder(_ root: TreeNode?) -> [[Int]] {
        guard let root = root else { return [] }
        var result: [[Int]] = []
        var level = [root]
        while !level.isEmpty {
            result.append(level.map { $0.val })
            var next: [TreeNode] = []
            for node in level {
                if let left = node.left { next.append(left) }
                if let right = node.right { next.append(right) }
            }
            level = next
        }
        return result
    }
}

計算量: O(N)。各ノードをちょうど 1 回処理します。