トポロジカルソートをPythonで実装して関数化する

hirohirohirohiros.hatenablog.com

概要

　トポロジカルソートは，有向無閉路グラフ(DAG)の各ノードについて，どのノードも出力先のノードより後ろに並ぶように並べるソートです．ようするに，ノードと一列に並べて，辺の向きが右向きのみになるように並び替えるソートのことです．(または左向きのみ)
　トポロジカルソートに掛けるグラフが，有向閉路グラフだったときこのソートは失敗するので，有向グラフが閉路かどうかの判定にも使えます．(こっちの用途の方が多そう？)

アルゴリズムの内容

　幅優先探索を使います．
　今回は，ノードを一列に並べて，辺の向きが右向きのみになるようにソートする事をイメージします．(a→b→c→...というイメージ)
　a→bという有向グラフがあったとき，aに一つ出力辺があり，bに一つ入力辺があると言うようにします．
　まず，各ノードについて，入力辺の本数を数えます．これは，a→bを表す入力a bが与えられたら，list[b] += 1とすれば良いです．これを全てのノードについて行い，入力辺の本数を要素としたリストを用意します．
　このリストについて，list[i] = 0であるようなノードiは入力辺が0ということです．つまり，そのノードは，x→iとなるようなノードxと辺が存在しないということなので，一番左に並べるのが適切であると分かります．よって，ソート済みのノードを集めたリストansに，list[i] = 0となるようなノードを全て入れます．(この順番は適当で良いです．トポロジカルソートはノードの順序さえ守られていれば良いので．)さらにノードi達について，それらをキューに入れます．
　つぎに，キューから一つノードiを取り出します．iの位置が確定したということで，iの出力辺先のノードjについてjの入力辺の数を-1します．そして，list[j] = 0となったら，入力辺が無くなったということでansに入れ，キューにも入れます．
　これをキューが空になるまで繰り返します．
　閉路があった時どうなるか考えます．a→b⇆cというような閉路があった時，まず入力辺0のaが取り出されます．そして，bの入力辺が-1されます．しかしbの入力辺の本数は0になりません．よってbはキューに入りません．そしてキューが空になるので処理が終了します．つまり，ansの中身はaのみになり，閉路分の2ノード足りません．
　このように閉路があるとansに入れられる事が無いので，閉路の数だけノード数が少なくなります．つまり，ソート済みリストのノードの数が，元のノードの数より少なかったら閉路，そうでなかったら無閉路と判定できます．
　計算量はO(V+E)となります．

pythonコード

　以上を実装しました．閉路の時-1を返すようにしてますが，こういう場合Noneを返す方が適切なんですかね……？(成功時返すのがリストでなく数値の時，-1と区別がつきにくいため)

from collections import deque
def topological_sort(node, input_edge):
    """トポロジカルソート

    有向グラフの順序を守るようにソートする
    閉路があるか判定も出来る
    計算量: O(E+V)

    Args:
        node (list): edge[i] = [a, b,...] iからa,b,...に辺が伸びている
        input_edge (list): input_edge[i] = iの入力辺の本数

    Returns:
        list or -1:閉路が存在しないとき
                      ソート済みのリスト
                   閉路が存在する時
                      -1   
    """
    N = len(input_edge)
    ans = [i for i in range(N) if input_edge[i] == 0]
    que = deque(ans)
    while que:
        q = que.popleft()
        for e in node[q]:
            input_edge[e] -= 1
            if input_edge[e] == 0:
                que.append(e)
                ans.append(e)
    if len(ans) == N:
        return ans
    return -1