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

29. Construct Binary Tree from Preorder and Inorder Traversal

読了目安 約2分

Tree / LC 105 (Medium) — 2 つの走査結果から元の木を復元する。

この章の目次

問題

LeetCode 105. Construct Binary Tree from Preorder and Inorder Traversal(Medium / カテゴリ: Tree)

同じ二分木を 2 通りに走査した結果が与えられます。 preorder は先行順(自分→左→右)、inorder は中間順(左→自分→右)の走査列です。 値はすべて相異なります。 元の木を復元して根を返します。

考え方

ヒント 1

preorder の先頭は必ず木全体の根です。 その値を inorder の中で探すと、何と何に分かれるかを考えてください。

ヒント 2

inorder で根より左にある要素が左の部分木、右にある要素が右の部分木です。 左の部分木の要素数が分かるので、preorder の続きも左用・右用に分割でき、再帰で復元できます。 値から inorder の添字を引く辞書(Dictionary)を先に作ると、根の位置が O(1) で分かります。

Swift 実装のポイント

  • 毎回 firstIndex(of:) で根を探すと O(N²) になります。[値: 添字] の Dictionary を前計算します。
  • 配列を切り出して渡す代わりに、inorder 上の範囲 (lo, hi) と「preorder の次に読む位置」の変数 1 つで管理すると、コピーなしで書けます。
  • preorder は「自分→左→右」の順なので、左の部分木を先に再帰で作れば読み取り位置は自然に進みます。
模範解答
Swift
class Solution {
    func buildTree(_ preorder: [Int], _ inorder: [Int]) -> TreeNode? {
        var indexOf: [Int: Int] = [:]
        for (i, value) in inorder.enumerated() {
            indexOf[value] = i
        }
        var next = 0
        func build(_ lo: Int, _ hi: Int) -> TreeNode? {
            if lo > hi { return nil }
            let value = preorder[next]
            next += 1
            let node = TreeNode(value)
            let mid = indexOf[value]!
            node.left = build(lo, mid - 1)
            node.right = build(mid + 1, hi)
            return node
        }
        return build(0, inorder.count - 1)
    }
}

計算量: O(N)。Dictionary の前計算の後は、ノード 1 個につき定数回の処理です。