fera_graph/algs/
kruskal.rs

1// This Source Code Form is subject to the terms of the Mozilla Public
2// License, v. 2.0. If a copy of the MPL was not distributed with this
3// file, You can obtain one at http://mozilla.org/MPL/2.0/.
4
5//! [Kruskal]'s minimum spanning tree algorithm.
6//!
7//! [Kruskal]: https://en.wikipedia.org/wiki/Kruskal's_algorithm
8
9use params::*;
10use prelude::*;
11use unionfind::{NewUnionFind, UnionFind, WithUnionFind};
12
13use fera_fun::vec;
14
15use std::ops::DerefMut;
16
17pub trait Visitor<G>
18where
19    G: WithEdge + WithUnionFind,
20{
21    fn accept(&mut self, g: &G, e: Edge<G>, ds: &mut UnionFind<G>) -> bool;
22
23    #[allow(unused_variables)]
24    fn after_union(&mut self, g: &G, e: Edge<G>, ds: &mut UnionFind<G>) {}
25}
26
27#[derive(Clone, Copy, Debug, PartialEq, Eq)]
28pub struct AcceptAll;
29
30impl<G> Visitor<G> for AcceptAll
31where
32    G: WithEdge + WithUnionFind,
33{
34    fn accept(&mut self, _g: &G, _e: Edge<G>, _ds: &mut UnionFind<G>) -> bool {
35        true
36    }
37}
38
39impl<F, G> Visitor<G> for F
40where
41    G: WithEdge + WithUnionFind,
42    F: FnMut(&G, Edge<G>, &mut UnionFind<G>) -> bool,
43{
44    fn accept(&mut self, g: &G, e: Edge<G>, ds: &mut UnionFind<G>) -> bool {
45        self(g, e, ds)
46    }
47}
48
49pub struct Iter<'a, G: 'a, E, V, U> {
50    g: &'a G,
51    edges: E,
52    visitor: V,
53    ds: U,
54}
55
56impl<'a, G, E, V, U> Iterator for Iter<'a, G, E, V, U>
57where
58    G: 'a + WithUnionFind,
59    E: Iterator,
60    E::Item: IntoOwned<Edge<G>>,
61    V: Visitor<G>,
62    U: DerefMut<Target = UnionFind<G>>,
63{
64    type Item = Edge<G>;
65
66    fn next(&mut self) -> Option<Edge<G>> {
67        if self.ds.num_sets() > 1 {
68            for e in self.edges.by_ref().map(|e| e.into_owned()) {
69                let (u, v) = self.g.ends(e);
70                if !self.ds.in_same_set(u, v) && self.visitor.accept(self.g, e, &mut self.ds) {
71                    self.ds.union(u, v);
72                    self.visitor.after_union(self.g, e, &mut self.ds);
73                    return Some(e);
74                }
75            }
76        }
77        None
78    }
79}
80
81pub trait Kruskal: WithUnionFind {
82    fn kruskal_mst<T, W>(
83        &self,
84        weight: W,
85    ) -> KruskalAlg<&Self, Vec<Edge<Self>>, AcceptAll, NewUnionFind<Self>>
86    where
87        W: EdgePropGet<Self, T>,
88        T: Ord,
89    {
90        self.kruskal().weight(weight)
91    }
92
93    fn kruskal(&self) -> KruskalAlg<&Self, AllEdges<Self>, AcceptAll, NewUnionFind<Self>> {
94        KruskalAlg(self, AllEdges(self), AcceptAll, NewUnionFind(self))
95    }
96}
97
98impl<G: WithUnionFind> Kruskal for G {}
99
100generic_struct! {
101    #[must_use]
102    pub struct KruskalAlg(graph, edges, visitor, unionfind)
103}
104
105impl<'a, G, E, V, U> KruskalAlg<&'a G, E, V, U>
106where
107    G: WithUnionFind,
108{
109    pub fn weight<W, T>(self, w: W) -> KruskalAlg<&'a G, Vec<Edge<G>>, V, U>
110    where
111        W: EdgePropGet<G, T>,
112        T: Ord,
113    {
114        let edges = vec(self.0.edges()).sorted_by_prop(&w);
115        self.edges(edges)
116    }
117}
118
119impl<'a, G, E, V, U> IntoIterator for KruskalAlg<&'a G, E, V, U>
120where
121    G: WithUnionFind,
122    E: IntoIterator,
123    E::Item: IntoOwned<Edge<G>>,
124    V: Visitor<G>,
125    U: ParamDerefMut<Target = UnionFind<G>>,
126{
127    type Item = Edge<G>;
128    type IntoIter = Iter<'a, G, E::IntoIter, V, U::Output>;
129
130    fn into_iter(self) -> Self::IntoIter {
131        let KruskalAlg(g, edges, visitor, ds) = self;
132        Iter {
133            g,
134            visitor,
135            edges: edges.into_iter(),
136            ds: ds.build(),
137        }
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::Kruskal;
144    use fera_fun::vec;
145    use prelude::*;
146
147    #[test]
148    fn kruskal_mst() {
149        let g: StaticGraph = graph!(
150            5,
151            (0, 4), // 0
152            (2, 3), // 1
153            (0, 1), // 2
154            (1, 4),
155            (1, 2), // 4
156            (2, 4),
157            (3, 4),
158        );
159        let mut weight = g.default_edge_prop(0usize);
160        for (e, &w) in g.edges().zip(&[1, 2, 3, 4, 5, 6, 7]) {
161            weight[e] = w;
162        }
163        let e = vec(g.edges());
164        assert_eq!(vec![e[0], e[1], e[2], e[4]], vec(g.kruskal_mst(&weight)));
165    }
166}
167
168// TODO: write benchmarks and optimize