二分木の幅優先探索をLevel Order Traversalで実現する

May 4th, 2021

Level order traversalとは

二分木が与えられた時に、同じ階層に存在するノードを全て走査してから次の階層を走査するようなアルゴリズムをLevel order traversal(レベル順走査)と言い、二分木に対する幅優先探索アルゴリズムとして知られている。。

例えば以下のような二分木が与えられたとき、Level order traversalでは[[1],[2,3],[4,5,6,7]]という順に走査される。

二分木


このような走査を実現するには、キューというFIFOのデータ構造を使うのが一般的であるので、まず初めにキューを使った解法をPythonコードを交えながら解説する。

なお、再帰関数を用いてもLevel Order Traversalを実現することができる。こちらは厳密には幅優先探索ではなく深さ優先探索になるのだが、こちらの手法でも同じ結果を出力する事ができるので、併せて紹介する。


Pythonで二分木を作る

アルゴリズムを実装する前に、まずは先ほどの二分木をPythonで実装する。

ノードの値、そしてleftとrightにポインタを持つノードクラスを定義し、ポインタを後続ノードと繋げれば良い。

level_order.py
class TreeNode:
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None

n1 = TreeNode(1)
n2 = TreeNode(2)
n3 = TreeNode(3)
n4 = TreeNode(4)
n5 = TreeNode(5)
n6 = TreeNode(6)
n7 = TreeNode(7)
n1.left = n2
n1.right = n3
n2.left = n4
n2.right = n5
n3.left = n6
n3.right = n7

キューとは

キューとは、FIFO(First In First Out)型のデータ構造で、文字通り最初に格納されたデータが最初に取り出される仕組みを持つ。

キューにデータを格納する事をエンキュー(enqueue)、データを取り出す事をデキュー(dequeue)と言い、Pythonではそれぞれ配列に対してappendpopで実現できる。

queue = [1,2,3]

# enqueue
queue.append(4)

# dequeue
queue.pop()

Level order Traversalの実装

キューを用いた解法

level order traversalのアルゴリズムは次のようになる。

  1. まず初めに、根ノードをキューに格納する。
  2. キューに格納されている全てのノードの値を空配列に追加する。
  3. その後、キューに格納されている全てのノードをその子ノードに更新する。子ノードが存在しない場合は何もしない。
  4. キューが空になるまで上記の処理を繰り返す。

上で紹介した二分木で言うと、まず初めに1という値が格納された根ノードがキューに格納される。

次に、キューに入っている全てのノード(現段階では根ノードのみ)の値を答えとなる配列に追加する。現段階で、配列の中身は[[1]]

その処理が終わった後に、キューに格納される全てのノード(現段階では根ノードのみ)の子ノード、つまり現段階では左の子ノードである2と右の子ノードである3、でキューを更新する。

この時注意しなければならないのは、根ノードが子を持たない時はNoneがキューに格納されてしまうので、その場合はNoneを削除する必要がある。

上記の作業をキューの中にノードが存在する限り繰り返し、キューの中身が全てNoneになったら処理は終了する。

ノードはキューに追加する1回とキューから取り出される1回の計2回の処理がNN回(N=ノードの数N=ノードの数)行われることになるが、オーダー記法では定数倍は無視されるので、計算量はO(N)O(N)となる。

*** 省略 ***

def levelOrder(root):
    ans, queue = [], [root]

    while len(queue) > 0:
        ans.append([node.val for node in queue])
        tmp = []
        for node in queue:
            tmp.extend([node.left, node.right])
            print(tmp)
        queue = [node for node in tmp if node]
    return ans

print(levelOrder(n1))

再帰関数を用いた深さ優先的解法

まず初めに、空配列ansと0を格納した変数indexを宣言し、levelOrder関数の引数として受け取った根ノードと併せて、新たに補助的に作成した関数(ここではhelperとする)に代入する。

ここで宣言したansは最終的な出力にあたり、indexは関数が実行される時点での木の深さを表す。

このhelper関数では、次のようなアルゴリズムを実行する。

  1. 引数として受け取ったノードがnullである場合は処理を終了する
  2. ansの中に格納されている配列の数よりもindex+1した数の方が大きければ、ansにから配列を追加する
  3. ans配列の0から数えて添字index番目にあたる配列に受け取ったノードの値を追加する。
  4. 受け取ったノードの両方の子ノードに対して同じ処理を繰り返す。

このアルゴリズムがなぜ深さ優先探索なのかは、上記アルゴリズムの④の部分を見れば分かる。

上で紹介した二分木を例にとって説明すると、まずhelper関数が最初に呼び出されてから③に至った段階では、ans配列の中身は[[1]]となっている。

次に、根ノードに対して2という値を持つ左の子ノードを引数に取ったhelper関数が呼び出されるのだが、この関数の再帰構造は、葉ノードにいきつくまでは終わらない。なぜなら、この関数の終了条件(基底部)は、引数として受け取ったノードがnullの場合のみだからだ。

よって、根ノードに大して右の子ノード(3という値を持つ)にhelper関数が実行される段階では、ans配列の中身は[[1], [2], [4, 5]]となっていることから、このアルゴリズムは深さ優先探索であることがわかる。

*** 省略 ***

def levelOrder(root):
    ans, index = [], 0
    helper(root, ans, index)
    print(ans)

def helper(root, ans, index):
    if not root:
        return

    if index + 1 > len(ans):
        ans.append([])
    ans[index].append(root.val)
    helper(root.left, ans, index+1)
    helper(root.right, ans, index+1)

print(levelOrder(n1))

参考