非再帰BFSでトポソから木DPをする

非再帰 BFS

Python で非再帰 BFS を書きます。ついでにトポロジカルソート（トポソ）もできます。トポロジカルソートを知らない人は適当に ググって ください。

トポロジカルソートを使うと、木 DP が簡単にできます。

目標

木の問題は基本的にトポソすればなんとかなります （本当？） 。みんな大好き 木DP もトポソ順に処理すればその都度 BFS / DFS する必要ないのでらくちんです。この記事では BFS で木のトポソを求めて、木 DP をやるのが目標です。

入力

まずは入力部分。1-indexed で辺リストを X に格納しています。

test.py

N = int(input()) # 頂点数
X = [[] for i in range(N)] # 辺リスト
for i in range(N-1):
    x, y = map(int, input().split())
    X[x-1].append(y-1)
    X[y-1].append(x-1)

BFS / トポロジカルソート

ここからトポソゾーン。
基本的には BFS です。別に DFS でもほとんどコード変わんないのでどっちでもいいけど。非再帰で BFS をするには deque を使うのが簡単です。順に各頂点を見て、その子リストを queue に追加します （通常、入力から作られるのは子リストではなく隣接頂点リストだと思います。この場合、隣接頂点リストのうち親でないものを子リストだと思えばよいです。また、追加のたびに、追加されたものの親を記録していきます） 。追加は右から、取得は左からすることによって BFS になります。

test.py

from collections import deque

P = [-1] * N # P[i] はiの親。iが根なら-1
Q = deque([0]) # queue。根にするやつを最初に追加
R = [] # トポロジカルソート
while Q:
    i = deque.popleft(Q)
    R.append(i)
    for a in X[i]:
        if a == P[i]: continue
        P[a] = i
        X[a].remove(i) # ☆☆☆
        deque.append(Q, a)

print("X =", X) # 子リスト
print("P =", P) # 親
print("R =", R) # トポロジカルソート

隣接頂点リスト から 子リスト に

ところで、隣接頂点リストより子リストになっている方が使いやすくないですか？私は使いやすい気がします。つまり、親に伸びる辺を削除したいです。これはコード中 ☆☆☆ のところで remove 関数で行っています。
勘のいい方は remove の計算量が気になるかもしれません。 Python の remove は O(要素数) なので2乗にならないか心配ですが、辺の数の合計は O(N) なので、全体の計算量は O(N) から変わりません。

でも remove はあまり行儀がよくない気がしますね。

木 DP

木 DP の定義はよく知らないのですが、根からまたは葉から順に決めていくような DP でしょうか。
具体的には、頂点 $i$ が $k$ 個の子 $j_0,\ ...,\ j_{k-1}$ を持つときに、 dp[i] が dp[j] たちから計算できるときにこれを使います。ここでは簡単のため dp[i] を 頂点 $i$ を根とする部分木のサイズ として木 DP をします。この場合、 dp[i] = 1 で初期化して、各 i の子 j について dp[i] += dp[j] をすればよいです。ただし更新の順番が大事で「葉」側から順に処理を行う必要があります。この場合、トポソを逆順に見ればよいです。簡単ですね。逆に「根」側から処理したいときはトポソを前から見ればよいです。

test.py

dp = [1] * N
for i in R[::-1]: # トポロジカルソートを逆順に（「葉」側から処理）
    for j in X[i]:
        dp[i] += dp[j]

全方位木 DP

気が向いたら全方位木 DP も書きます。

2020/04/08 追記
書きました