配列から木構造を組み立てる【TypeScript】

はじめに

こんにちは！Shungiku です。
最近は就活のために専ら LeetCode の問題を解いていることが多いです。

さて、LeetCode では木構造の問題（例えばこれとか）において、入力のサンプルが以下のような配列で与えられます。

root = [3,4,5,1,2,null,null,null,null,0]

一方、実際に Submit した場合、関数（またはメソッド）の引数で受け取る入力値は木構造のルートのノードになっています。

ちなみにノードの構造は LeetCode では以下のように定義されています。

class TreeNode {
    val: number
    left: TreeNode | null = null
    right: TreeNode | null = null

    constructor(val: number) {
        this.val = val
    }
}

「入力サンプルの配列をTreeNodeを用いた木構造に変換してくれるヘルパー関数がほしい！」というのが、本稿のモチベーションです。

そんなヘルパー関数があれば、Failしたケースについてデバッグする際、サクッと入力値を用意できて便利です。

実装

配列から木構造を組み立てる関数は以下のようになります。

function buildTree(array: (number | null)[]): TreeNode | null {
    if (array.length === 0 || array[0] === null) {
        return null
    }

    const root = new TreeNode(array[0])

    const queue: (TreeNode | null)[] = [root]

    for (let i=1; i < array.length; i++) {
        const val = array[i]
        const node = val === null ? null : new TreeNode(val)
        queue.push(node) // ノードを作成したら必ず `queue` に突っ込む
        if (i % 2 === 0) { // right
            if (val !== null && queue[0] !== null) {
                queue[0].right = node // `queue` の先頭のノードを親ノードとして扱う
            }
            queue.shift() // right 側の処理が終わったら、それまでの親ノードは用無しなので `queue` から追い出す
        }
        else { // left
            if (val !== null && queue[0] !== null) {
                queue[0].left = node // `queue` の先頭のノードを親ノードとして扱う
            }
        }
    }

    return root
}

「新たに生成した子ノードをどの親ノードにくっつけるか」は queue で管理しています。

ポイントは以下の通りです。

• ノードを作成したら必ず queue に突っ込む
• queue の先頭のノードを親ノードとして扱う
• 配列のインデックスが奇数の場合は left 、偶数の場合は right 側に子ノードをくっつける
• 親ノードまたは子ノードが null の場合は
• right 側の処理が終わったら、それまでの親ノードは用無しなので queue から追い出す

まとめ

これで Fail したケースに対してのデバッグが捗るぜ...!

ちなみに連結リストについても同様の記事を執筆してますので、以下ご参照ください。