Algod/graph/
prim.rs

1use std::cmp::Reverse;
2use std::collections::{BTreeMap, BinaryHeap};
3use std::ops::Add;
4
5type Graph<V, E> = BTreeMap<V, BTreeMap<V, E>>;
6
7fn add_edge<V: Ord + Copy, E: Ord + Add + Copy>(graph: &mut Graph<V, E>, v1: V, v2: V, c: E) {
8    graph.entry(v1).or_insert_with(BTreeMap::new).insert(v2, c);
9    graph.entry(v2).or_insert_with(BTreeMap::new).insert(v1, c);
10}
11
12// selects a start and run the algorithm from it
13pub fn prim<V: Ord + Copy + std::fmt::Debug, E: Ord + Add + Copy + std::fmt::Debug>(
14    graph: &Graph<V, E>,
15) -> Graph<V, E> {
16    match graph.keys().next() {
17        Some(v) => prim_with_start(graph, *v),
18        None => BTreeMap::new(),
19    }
20}
21
22// only works for a connected graph
23// if the given graph is not connected it will return the MST of the connected subgraph
24pub fn prim_with_start<V: Ord + Copy, E: Ord + Add + Copy>(
25    graph: &Graph<V, E>,
26    start: V,
27) -> Graph<V, E> {
28    // will contain the MST
29    let mut mst: Graph<V, E> = Graph::new();
30    // a priority queue based on a binary heap, used to get the cheapest edge
31    // the elements are an edge: the cost, destination and source
32    let mut prio = BinaryHeap::new();
33
34    mst.insert(start, BTreeMap::new());
35
36    for (v, c) in &graph[&start] {
37        // the heap is a max heap, we have to use Reverse when adding to simulate a min heap
38        prio.push(Reverse((*c, v, start)));
39    }
40
41    while let Some(Reverse((dist, t, prev))) = prio.pop() {
42        // the destination of the edge has already been seen
43        if mst.contains_key(t) {
44            continue;
45        }
46
47        // the destination is a new vertex
48        add_edge(&mut mst, prev, *t, dist);
49
50        for (v, c) in &graph[t] {
51            if !mst.contains_key(v) {
52                prio.push(Reverse((*c, v, *t)));
53            }
54        }
55    }
56
57    mst
58}
59
60#[cfg(test)]
61mod tests {
62    use super::{add_edge, prim, Graph};
63    use std::collections::BTreeMap;
64
65    #[test]
66    fn empty() {
67        assert_eq!(prim::<usize, usize>(&BTreeMap::new()), BTreeMap::new());
68    }
69
70    #[test]
71    fn single_vertex() {
72        let mut graph: Graph<usize, usize> = BTreeMap::new();
73        graph.insert(42, BTreeMap::new());
74
75        assert_eq!(prim(&graph), graph);
76    }
77
78    #[test]
79    fn single_edge() {
80        let mut graph = BTreeMap::new();
81
82        add_edge(&mut graph, 42, 666, 12);
83
84        assert_eq!(prim(&graph), graph);
85    }
86
87    #[test]
88    fn tree_1() {
89        let mut graph = BTreeMap::new();
90
91        add_edge(&mut graph, 0, 1, 10);
92        add_edge(&mut graph, 0, 2, 11);
93        add_edge(&mut graph, 2, 3, 12);
94        add_edge(&mut graph, 2, 4, 13);
95        add_edge(&mut graph, 1, 5, 14);
96        add_edge(&mut graph, 1, 6, 15);
97        add_edge(&mut graph, 3, 7, 16);
98
99        assert_eq!(prim(&graph), graph);
100    }
101
102    #[test]
103    fn tree_2() {
104        let mut graph = BTreeMap::new();
105
106        add_edge(&mut graph, 1, 2, 11);
107        add_edge(&mut graph, 2, 3, 12);
108        add_edge(&mut graph, 2, 4, 13);
109        add_edge(&mut graph, 4, 5, 14);
110        add_edge(&mut graph, 4, 6, 15);
111        add_edge(&mut graph, 6, 7, 16);
112
113        assert_eq!(prim(&graph), graph);
114    }
115
116    #[test]
117    fn tree_3() {
118        let mut graph = BTreeMap::new();
119
120        for i in 1..100 {
121            add_edge(&mut graph, i, 2 * i, i);
122            add_edge(&mut graph, i, 2 * i + 1, -i);
123        }
124
125        assert_eq!(prim(&graph), graph);
126    }
127
128    #[test]
129    fn graph_1() {
130        let mut graph = BTreeMap::new();
131        add_edge(&mut graph, 'a', 'b', 6);
132        add_edge(&mut graph, 'a', 'c', 7);
133        add_edge(&mut graph, 'a', 'e', 2);
134        add_edge(&mut graph, 'a', 'f', 3);
135        add_edge(&mut graph, 'b', 'c', 5);
136        add_edge(&mut graph, 'c', 'e', 5);
137        add_edge(&mut graph, 'd', 'e', 4);
138        add_edge(&mut graph, 'd', 'f', 1);
139        add_edge(&mut graph, 'e', 'f', 2);
140
141        let mut ans = BTreeMap::new();
142        add_edge(&mut ans, 'd', 'f', 1);
143        add_edge(&mut ans, 'e', 'f', 2);
144        add_edge(&mut ans, 'a', 'e', 2);
145        add_edge(&mut ans, 'b', 'c', 5);
146        add_edge(&mut ans, 'c', 'e', 5);
147
148        assert_eq!(prim(&graph), ans);
149    }
150
151    #[test]
152    fn graph_2() {
153        let mut graph = BTreeMap::new();
154        add_edge(&mut graph, 1, 2, 6);
155        add_edge(&mut graph, 1, 3, 1);
156        add_edge(&mut graph, 1, 4, 5);
157        add_edge(&mut graph, 2, 3, 5);
158        add_edge(&mut graph, 2, 5, 3);
159        add_edge(&mut graph, 3, 4, 5);
160        add_edge(&mut graph, 3, 5, 6);
161        add_edge(&mut graph, 3, 6, 4);
162        add_edge(&mut graph, 4, 6, 2);
163        add_edge(&mut graph, 5, 6, 6);
164
165        let mut ans = BTreeMap::new();
166        add_edge(&mut ans, 1, 3, 1);
167        add_edge(&mut ans, 4, 6, 2);
168        add_edge(&mut ans, 2, 5, 3);
169        add_edge(&mut ans, 2, 3, 5);
170        add_edge(&mut ans, 3, 6, 4);
171
172        assert_eq!(prim(&graph), ans);
173    }
174
175    #[test]
176    fn graph_3() {
177        let mut graph = BTreeMap::new();
178        add_edge(&mut graph, "v1", "v2", 1);
179        add_edge(&mut graph, "v1", "v3", 3);
180        add_edge(&mut graph, "v1", "v5", 6);
181        add_edge(&mut graph, "v2", "v3", 2);
182        add_edge(&mut graph, "v2", "v4", 3);
183        add_edge(&mut graph, "v2", "v5", 5);
184        add_edge(&mut graph, "v3", "v4", 5);
185        add_edge(&mut graph, "v3", "v6", 2);
186        add_edge(&mut graph, "v4", "v5", 2);
187        add_edge(&mut graph, "v4", "v6", 4);
188        add_edge(&mut graph, "v5", "v6", 1);
189
190        let mut ans = BTreeMap::new();
191        add_edge(&mut ans, "v1", "v2", 1);
192        add_edge(&mut ans, "v5", "v6", 1);
193        add_edge(&mut ans, "v2", "v3", 2);
194        add_edge(&mut ans, "v3", "v6", 2);
195        add_edge(&mut ans, "v4", "v5", 2);
196
197        assert_eq!(prim(&graph), ans);
198    }
199}
Algod/graph/prim.rs

Algod/graph/
prim.rs
