非再帰DFS

PyPy では深い再帰が遅く、スタックあふれも起きやすいので、スタックで DFS を書きます。帰りがけのタイミングを知るために、頂点 $v$ の訪問時にスタックへ ~v（ $= -v-1$）を積んでおき、それを取り出したときを帰りがけとします。

• dfs(g, s): 隣接リスト g[v] = [to, ...] で $s$ から DFS し、pre（行きがけ順）、post（帰りがけ順）、par（DFS 木での親、 $s$ と未到達の頂点は $-1$）を返す
• 隣接リストの順に頂点を訪れるので、再帰で書いた DFS と同じ順序になる
• 計算量: $O(N + M)$

コード中の「行きがけの処理」「帰りがけの処理」の位置に処理を書き足して使うこともできます。木 DP なら、帰りがけの処理で par[v] に $v$ の値を集めます。

コード

def dfs(g, s):
    n = len(g)
    seen = [False] * n
    par = [-1] * n
    pre = []
    post = []
    st = [s]
    while st:
        v = st.pop()
        if v >= 0:
            if seen[v]:
                continue
            seen[v] = True
            pre.append(v)
            # 行きがけの処理
            st.append(~v)
            for to in reversed(g[v]):
                if not seen[to]:
                    par[to] = v
                    st.append(to)
        else:
            v = ~v
            post.append(v)
            # 帰りがけの処理
    return pre, post, par