prepona/algo/mst/
kruskal.rs

1use std::cell::RefCell;
2use std::collections::HashSet;
3use std::rc::Rc;
4
5use crate::graph::{subgraph::Subgraph, Edge, UndirectedEdge};
6use crate::provide;
7
8/// Finds minimum spanning tree using kruskal algorithm.
9///
10/// # Examples
11/// ```
12/// use prepona::prelude::*;
13/// use prepona::storage::Mat;
14/// use prepona::graph::MatGraph;
15/// use prepona::algo::Kruskal;
16///
17/// //  Given: Graph
18/// //                5
19/// //      f ----------------.
20/// //      |                 |
21/// //    3 |  1     1     4  |
22/// //      a --- b --- d --- e
23/// //    3 |   5 |   2 |   1 |
24/// //      |     |     |     |
25/// //      c ----'-----'-----'
26/// let mut graph = MatGraph::init(Mat::<usize>::init());
27/// let a = graph.add_vertex();
28/// let b = graph.add_vertex();
29/// let c = graph.add_vertex();
30/// let d = graph.add_vertex();
31/// let e = graph.add_vertex();
32/// let f = graph.add_vertex();
33///
34/// let ab = graph.add_edge_unchecked(a, b, 1.into());
35/// graph.add_edge_unchecked(a, c, 3.into());
36/// let af = graph.add_edge_unchecked(a, f, 3.into());
37///
38/// graph.add_edge_unchecked(b, c, 5.into());
39/// let bd = graph.add_edge_unchecked(b, d, 1.into());
40///
41/// let dc = graph.add_edge_unchecked(d, c, 2.into());
42/// graph.add_edge_unchecked(d, e, 4.into());
43///
44/// let ec = graph.add_edge_unchecked(e, c, 1.into());
45/// graph.add_edge_unchecked(e, f, 5.into());
46///
47/// let mut tags = std::collections::HashMap::<usize, &'static str>::new();
48/// tags.insert(a, "a");
49/// tags.insert(b, "b");
50/// tags.insert(c, "c");
51/// tags.insert(d, "d");
52/// tags.insert(e, "e");
53/// tags.insert(f, "f");
54///
55/// let mst = Kruskal::init(&graph).execute(&graph);
56///
57/// assert_eq!(mst.vertex_count(), 6);
58/// assert_eq!(mst.edges_count(), 5);
59/// assert!(vec![ab, af, bd, dc, ec].into_iter().all(|edge_id| mst.edge(edge_id).is_ok()))
60/// ```
61pub struct Kruskal {
62    sets: Vec<Rc<RefCell<HashSet<usize>>>>,
63}
64
65impl Kruskal {
66    /// Initializes the structure.
67    pub fn init<G, W: Ord, E: Edge<W>>(graph: &G) -> Self
68    where
69        G: provide::Vertices + provide::Edges<W, E> + provide::Graph<W, E, UndirectedEdge>,
70    {
71        let vertex_count = graph.vertex_count();
72
73        let mut sets = vec![];
74        sets.resize_with(vertex_count, || Rc::new(RefCell::new(HashSet::new())));
75
76        for virt_id in 0..vertex_count {
77            sets[virt_id].borrow_mut().insert(virt_id);
78        }
79
80        Kruskal { sets }
81    }
82
83    /// Finds minimum spanning tree.
84    ///
85    /// # Arguments
86    /// `graph`: Graph to find its MST.
87    ///
88    /// # Returns
89    /// MST as a subgraph of the original graph(containing vertices and selected edges)
90    pub fn execute<'a, G, W: Ord, E: Edge<W>>(
91        mut self,
92        graph: &'a G,
93    ) -> Subgraph<W, E, UndirectedEdge, G>
94    where
95        G: provide::Edges<W, E>
96            + provide::Neighbors
97            + provide::Vertices
98            + provide::Graph<W, E, UndirectedEdge>,
99    {
100        let mut mst = Vec::<(usize, usize, usize)>::new();
101
102        let id_map = graph.continuos_id_map();
103
104        let mut edges = graph.edges();
105
106        edges.sort_by(|(_, _, e1), (_, _, e2)| e1.get_weight().cmp(e2.get_weight()));
107
108        for (v_real_id, u_real_id, edge) in edges {
109            let v_virt_id = id_map.virt_id_of(v_real_id);
110            let u_virt_id = id_map.virt_id_of(u_real_id);
111
112            if !self.sets[v_virt_id]
113                .borrow()
114                .eq(&*self.sets[u_virt_id].borrow())
115            {
116                mst.push((v_real_id, u_real_id, edge.get_id()));
117
118                let union_set = self.sets[v_virt_id]
119                    .borrow()
120                    .union(&*self.sets[u_virt_id].borrow())
121                    .copied()
122                    .collect::<HashSet<usize>>();
123
124                let sharable_set = Rc::new(RefCell::new(union_set));
125
126                for member in sharable_set.borrow().iter() {
127                    self.sets[*member] = sharable_set.clone();
128                }
129            }
130        }
131
132        let vertices = mst
133            .iter()
134            .flat_map(|(src_id, dst_id, _)| vec![*src_id, *dst_id])
135            .collect::<HashSet<usize>>();
136
137        Subgraph::init(graph, mst, vertices)
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144    use crate::graph::MatGraph;
145    use crate::provide::*;
146    use crate::storage::Mat;
147
148    #[test]
149    fn empty_graph() {
150        let graph = MatGraph::init(Mat::<usize>::init());
151
152        let mst = Kruskal::init(&graph).execute(&graph);
153
154        assert_eq!(mst.vertex_count(), 0);
155    }
156
157    #[test]
158    fn trivial_directed_graph() {
159        //  Given: Graph
160        //                5
161        //      f ----------------.
162        //      |                 |
163        //    3 |  1     1     4  |
164        //      a --- b --- d --- e
165        //    3 |   5 |   2 |   1 |
166        //      |     |     |     |
167        //      c ----'-----'-----'
168        let mut graph = MatGraph::init(Mat::<usize>::init());
169        let a = graph.add_vertex();
170        let b = graph.add_vertex();
171        let c = graph.add_vertex();
172        let d = graph.add_vertex();
173        let e = graph.add_vertex();
174        let f = graph.add_vertex();
175
176        let ab = graph.add_edge_unchecked(a, b, 1.into());
177        graph.add_edge_unchecked(a, c, 3.into());
178        let af = graph.add_edge_unchecked(a, f, 3.into());
179
180        graph.add_edge_unchecked(b, c, 5.into());
181        let bd = graph.add_edge_unchecked(b, d, 1.into());
182
183        let dc = graph.add_edge_unchecked(d, c, 2.into());
184        graph.add_edge_unchecked(d, e, 4.into());
185
186        let ec = graph.add_edge_unchecked(e, c, 1.into());
187        graph.add_edge_unchecked(e, f, 5.into());
188
189        let mut tags = std::collections::HashMap::<usize, &'static str>::new();
190        tags.insert(a, "a");
191        tags.insert(b, "b");
192        tags.insert(c, "c");
193        tags.insert(d, "d");
194        tags.insert(e, "e");
195        tags.insert(f, "f");
196
197        let mst = Kruskal::init(&graph).execute(&graph);
198
199        assert_eq!(mst.vertex_count(), 6);
200        assert_eq!(mst.edges_count(), 5);
201        assert!(vec![ab, af, bd, dc, ec]
202            .into_iter()
203            .all(|edge_id| mst.edge(edge_id).is_ok()))
204    }
205}