Rustコトハジメ

プログラミング言語Rustに関する情報をお届けします。

Chu-Liu/Edmondsアルゴリズムを強連結成分分解を使って実装した

www.rustforbeginners.com

で、SCCについて説明しました。今回は、SCCを応用して、最小全域有向木を探すアルゴリズムを実装します。一番参考になったのはこちらです。アルゴリズムとしては同じですが、SCCを使って自分で実装してみます。

www.creativ.xyz

アルゴリズムの概要は以下です。

  1. 現状のグラフにおいて各点(ルート以外)で、入力エッジのうちコスト最小のものを選択する
  2. それが木になってるか確かめて、もしそうであれば、明らかにこれがコスト最小なので採用しておわり
  3. 木になっていない場合、サイクルがあってグラフが分断していることを意味している。これはSCCを計算することでわかる。点が1つより多い強連結成分に属していれば、サイクルしてることになる
  4. サイクルしている部分のコスト自体は採用してしまい、サイクルを縮約した上で(ここがうまい)、「1でもし最小以外のものを選んでいたら」を考える。すでに最小コスト分は採用してしまっているので、これらのエッジのコストは、1で採用したエッジとの差分とする
  5. 縮約したグラフに対して再帰して追加コストを計算する
mod chu_liu_edmonds {
    use crate::graph::scc;
    #[derive(Debug, Clone)]
    struct Edge(usize, u64);
    fn min_edge(edges: &[Edge]) -> &Edge {
        let mut r = &edges[0];
        for e in edges {
            if e.1 < r.1 {
                r = e;
            }
        }
        r
    }
    static NULL_EDGE: &'static Edge = &Edge(1<<40, 0);
    fn chu_liu_edmonds(in_g: &[Vec<Edge>], root: usize) -> u64 {
        let mut min_in_g: Vec<&Edge> = vec![];
        let mut min_out_g: Vec<Vec<usize>> = vec![vec![]; in_g.len()];
        for to in 0..in_g.len() {
            if to == root {
                min_in_g.push(NULL_EDGE);
                continue;
            }
            let e = min_edge(&in_g[to]);
            min_in_g.push(e);
            min_out_g[e.0].push(to);
        }

        let mut scc = scc::SCC::new(&min_out_g);
        scc.build();

        let mut max_cmp = 0;
        for &cmp in &scc.order {
            if cmp > max_cmp {
                max_cmp = cmp;
            }
        }

        let no_loop = max_cmp == scc.order.len()-1;
        if no_loop {
            let mut res = 0;
            for e in &min_in_g {
                res += e.1;
            }
            return res;
        }

        let mut groups = vec![vec![]; max_cmp+1];
        for v in 0..scc.order.len() {
            let cmp = scc.order[v];
            groups[cmp].push(v);
        }

        let mut contracted_cost = 0;
        let mut new_in_g = vec![vec![]; max_cmp+1];
        for group in groups {
            if group.len() > 1 { // loop
                let cmp_to = scc.order[group[0]];
                for &v in &group {
                    let cur_e = min_in_g[v];

                    contracted_cost += cur_e.1;

                    for e in &in_g[v] {
                        let in_group = group.contains(&e.0);
                        if !in_group {
                            let cmp_from = scc.order[e.0];
                            let diff_cost = e.1 - cur_e.1;
                            new_in_g[cmp_to].push(Edge(cmp_from, diff_cost));
                        }
                    }
                }
            } else {
                assert!(group.len() == 1);
                let v = group[0];
                for e in &in_g[v] {
                    let cmp_to = scc.order[v];
                    let cmp_from = scc.order[e.0];
                    new_in_g[cmp_to].push(Edge(cmp_from, e.1));
                }
            }
        }

        let new_root = scc.order[root];

        contracted_cost + chu_liu_edmonds(&new_in_g, new_root)
    }
}