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

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

pathのノード情報をリストでほしい場合

再帰

gist.github.com

を実行すると

>> [[0, 1, 3], [0, 1, 4], [0, 2, 5], [0, 2, 6]]

が出力されます.

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

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

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

print(list(path_generator(root)))

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

スタック

def path_generator(node):
    stack = [[node]]
    while stack:
        node = stack[-1][-1]
        if node is leaf:
            yield stack.pop()
        else:
            path = stack.pop()
            for child in children of node:
                stack.append(path + [child])

条件に当てはまるパスの個数を知りたい場合

leetcodeなどで使うのはこっちのが多いですかね。やっぱリストの結合をバンバンやるのは時間がかかるので。

rootからleafへのpathに含まれるノードの値のカウント

counter を bitmapに変えるなどのバリエーションがあります。

https://leetcode.com/problems/pseudo-palindromic-paths-in-a-binary-tree/

def dfs(node, conter):
    ans = 0
    counter += collections.Counter([node.val])
    if node.left is None and node.right is None:
        # reach leaf node
        # check something and update ans
    elif node.left is not None:
        ans += dfs(node.left, counter)
    elif node.right is not None:
        ans += dfs(node.right, counter)
    counter -= collections.Counter([node.val])
    return ans