haar_lib/graph/
prim.rs

1//! 最小全域木 (Prim)
2
3use crate::graph::*;
4use std::{cmp::Reverse, collections::BinaryHeap};
5
6/// Prim法
7///
8/// グラフが連結ならばSomeに包んで最小全域木の辺集合を返す。
9/// 非連結ならばNoneを返す。
10pub fn prim<E: EdgeTrait>(g: &Graph<Undirected, E>) -> Option<Vec<&E>>
11where
12    E::Weight: Ord,
13{
14    let n = g.len();
15    let mut visit = vec![false; n];
16    let mut ret = vec![];
17    let mut heap = BinaryHeap::new();
18
19    visit[0] = true;
20    for (index, e) in g.nodes[0].edges.iter().enumerate() {
21        heap.push(Reverse((e.weight(), e.from(), index)));
22    }
23
24    while let Some(Reverse((_, from, index))) = heap.pop() {
25        let e = &g.nodes[from].edges[index];
26        if visit[e.from()] == visit[e.to()] {
27            continue;
28        }
29
30        let i = if visit[e.from()] { e.to() } else { e.from() };
31        for (index, e) in g.nodes[i].edges.iter().enumerate() {
32            heap.push(Reverse((e.weight(), e.from(), index)));
33        }
34
35        visit[i] = true;
36
37        ret.push(e);
38    }
39
40    (ret.len() == n - 1).then_some(ret)
41}
42
43#[cfg(test)]
44mod tests {
45    use super::*;
46
47    #[test]
48    fn test() {
49        let mut g = Graph::<Undirected, _>::new(6);
50        g.extend(
51            vec![
52                (0, 1, 1),
53                (0, 2, 3),
54                (1, 2, 1),
55                (1, 3, 7),
56                (2, 4, 1),
57                (1, 4, 3),
58                (3, 4, 1),
59                (3, 5, 1),
60                (4, 5, 6),
61            ]
62            .into_iter()
63            .map(|(u, v, w)| Edge::new(u, v, w, ())),
64        );
65
66        let ans = prim(&g).unwrap().iter().map(|e| e.weight).sum::<i32>();
67
68        assert_eq!(ans, 5);
69    }
70}
haar_lib/graph/prim.rs

haar_lib/graph/
prim.rs
