pathfinding/undirected/
kruskal.rs

1//! Find minimum-spanning-tree in an undirected graph using
2//! [Kruskal's algorithm](https://en.wikipedia.org/wiki/Kruskal's_algorithm).
3
4use crate::FxIndexSet;
5use std::hash::Hash;
6use std::mem;
7
8// Find parent and compress path by path halving.
9fn find(parents: &mut [usize], mut node: usize) -> usize {
10    while parents[node] != node {
11        parents[node] = parents[parents[node]];
12        node = parents[node];
13    }
14    node
15}
16
17fn union(parents: &mut [usize], ranks: &mut [usize], mut a: usize, mut b: usize) {
18    if ranks[a] < ranks[b] {
19        mem::swap(&mut a, &mut b);
20    }
21    parents[b] = a;
22    if ranks[a] == ranks[b] {
23        ranks[a] += 1;
24    }
25}
26
27/// Minimal-spanning-tree for nodes with integer indices. The nodes must have
28/// consecutives indices between 0 and `number_of_nodes`-1.
29///
30/// # Panics
31///
32/// This function panics if a node is outside the range [0, `number_of_nodes`-1].
33pub fn kruskal_indices<C>(
34    number_of_nodes: usize,
35    edges: impl AsRef<[(usize, usize, C)]>,
36) -> impl Iterator<Item = (usize, usize, C)>
37where
38    C: Clone + Ord,
39{
40    let mut parents = (0..number_of_nodes).collect::<Vec<_>>();
41    let mut ranks = vec![1; number_of_nodes];
42    let mut edges = edges.as_ref().to_vec();
43    edges.sort_unstable_by(|a, b| a.2.cmp(&b.2));
44    edges.into_iter().filter_map(move |(a, b, w)| {
45        let ra = find(&mut parents, a);
46        let rb = find(&mut parents, b);
47        if ra == rb {
48            None
49        } else {
50            union(&mut parents, &mut ranks, ra, rb);
51            Some((a, b, w))
52        }
53    })
54}
55
56/// Find a minimum-spanning-tree. From a collection of
57/// weighted edges, return a vector of edges forming
58/// a minimum-spanning-tree.
59pub fn kruskal<N, C>(edges: &[(N, N, C)]) -> impl Iterator<Item = (&N, &N, C)>
60where
61    N: Hash + Eq,
62    C: Clone + Ord,
63{
64    let mut nodes = FxIndexSet::default();
65    let edges = edges
66        .iter()
67        .map(|(a, b, w)| {
68            let ia = nodes.insert_full(a).0;
69            let ib = nodes.insert_full(b).0;
70            (ia, ib, w.clone())
71        })
72        .collect::<Vec<_>>();
73    kruskal_indices(nodes.len(), edges).filter_map(move |(ia, ib, w)| {
74        Some((
75            <&N>::clone(nodes.get_index(ia)?), // Cannot fail
76            <&N>::clone(nodes.get_index(ib)?), // Cannot fail
77            w,
78        ))
79    })
80}
81
82#[cfg(test)]
83mod tests {
84    use super::find;
85
86    #[test]
87    fn path_halving() {
88        let mut parents = vec![0, 0, 1, 2, 3, 4, 5, 6];
89        assert_eq!(find(&mut parents, 7), 0);
90        assert_eq!(parents, vec![0, 0, 1, 1, 3, 3, 5, 5]);
91        assert_eq!(find(&mut parents, 7), 0);
92        assert_eq!(parents, vec![0, 0, 1, 0, 3, 3, 5, 3]);
93        assert_eq!(find(&mut parents, 7), 0);
94        assert_eq!(parents, vec![0, 0, 1, 0, 3, 3, 5, 0]);
95        assert_eq!(find(&mut parents, 6), 0);
96        assert_eq!(parents, vec![0, 0, 1, 0, 3, 3, 3, 0]);
97        assert_eq!(find(&mut parents, 6), 0);
98        assert_eq!(parents, vec![0, 0, 1, 0, 3, 3, 0, 0]);
99    }
100}
pathfinding/undirected/kruskal.rs

pathfinding/undirected/
kruskal.rs
