23日目: TopologicalSortを実装する

競プロ用ライブラリを作る Advent Calendar 2024の23日目です.

TopologicalSortって?

DAG (有向非巡回グラフ) の頂点を良い感じに並べ替えるアルゴリズムです. DAGとはループが無い有向グラフのことです

↓はA→C→D→Aと辿ると元の場所に戻ってこれるのでDAGではないです.

↓はどう矢印を辿っても元の場所に戻ってこれないのでDAGです

これを[A, C, B, D, E, F]のように, 矢印の行先が必ず後ろの要素になるように並べ替えるのがTopologicalSortです.
[A, B, C, E, D, F]みたいな順番でも条件を満たしますが, こういうときはどちらを返してもいいものとします.

仕組み

仕組みは簡単,「追加しても良い頂点を追加する」を繰り返すだけです.
ここで追加しても良い頂点は「その頂点に向かって辺が向かってる頂点が全部既に配列に入ってる頂点」です.

具体的にアルゴリズムを説明します

まず, 最初に「各頂点に入ってくる辺はいくつあるか」を数えておきます.
そしてそこから「入ってくる辺が無い頂点」をスタックに入れます. このスタックは追加しても良い頂点が入るスタックです.

そして

1. スタックから要素を取り出す
2. その要素を答えの配列の末尾に追加する
3. その頂点から出てる辺全部について以下のことを行う
   1. その辺が入る頂点の「入ってくる辺のカウント」を1つ減らす
   2. カウントが0になったなら, その頂点をスタックに入れる

を繰り返せばTopologicalSortできます.

pub fn sample_topological_sort(edges: &[Vec<usize>]) -> Vec<usize> {
    let n = edges.len();
    // 各頂点に入る辺の数を数える
    let mut count = std::iter::repeat_n(0usize, n).collect::<Vec<_>>();
    for v in edges {
        for &u in v {
            count[u] += 1;
        }
    }
    // 入ってくる頂点が無い頂点をスタックに入れる
    let mut stack = count
        .iter()
        .enumerate()
        .filter_map(|(i, &c)| (c == 0).then_some(i))
        .collect::<Vec<_>>();

    let mut ret = vec![];
    // スタックに入ってる頂点が無くなるまで取り出して処理する
    while let Some(v) = stack.pop() {
        // 結果の配列に頂点を入れる
        ret.push(v);
        // その頂点から出てる辺でループを回す
        for &u in &edges[v] {
            // 入る頂点のカウントを減らす
            count[u] -= 1;
            // カウントが0になったらその頂点をスタックに入れる
            if count[u] == 0 {
                stack.push(u);
            }
        }
    }
    ret
}

実装

• topologicalsort.rs

やるだけです. 柔軟性を持たせるためにイテレータを返すようにしてみました.

まとめ

いかがでしたか?

明日はAtCoder出禁データ構造を作ります.