サブロウ丸

主にプログラミングと数学

【Python】generatorを用いた, 二分木におけるパスの全列挙

このような番号付けで表現される二分木におけるpathの全列挙を考えます.

f:id:inarizuuuushi:20191020001106p:plain:w400

gist.github.com

を実行すると

>> [[0, 1, 3, 7], [0, 1, 3, 8], [0, 1, 4, 9], [0, 1, 4, 10], [0, 2, 5, 11], [0, 2, 5, 12], [0, 2, 6, 13], [0, 2, 6, 14]]

が出力されます.

深さ優先探索で左側からノードを探索していき, 葉ノードまでたどり着くと, _dfs()再帰関数のif文に入ります.
ここに入ったあとは, 自分のノードをリストの先頭に加えながら再帰関数を逆戻りするという具合ですね. generatorなので, list()などでラップしてやれば, pathのリストを得られます.



この実装は2分木でなくても, 一般の木に対しても応用することができます.

def dfs():
    return list(_dfs(0))

def _dfs(node):
    if node is leaf:
        yield [node]
    else:
        for child in children of node:
            for path in _dfs(child):
                yield [node] + path

print(dfs())

擬似コードで書くとこんな感じですね.