1use std::collections::HashMap;
2
3use derivative::Derivative;
4
5use crate::node::{Node, Position};
6use crate::union_find::UnionFind;
7
8#[derive(Derivative)]
9#[derivative(Clone, PartialEq, Eq, Hash)]
10pub struct Edge {
11 index: usize,
12 pub source: usize,
13 pub destination: usize,
14 #[derivative(PartialEq = "ignore")]
15 #[derivative(Hash = "ignore")]
16 pub weight: f32,
17}
18
19impl Edge {
20 pub fn from(index: usize, source: usize, destination: usize, weight: f32) -> Edge {
21 return Edge {
22 index,
23 source,
24 destination,
25 weight,
26 };
27 }
28}
29
30pub struct Graph {
31 pub edges_lookup: HashMap<usize, Edge>,
32 pub nodes_lookup: HashMap<usize, Node>,
33 pub node_position_lookup: Option<HashMap<usize, Position>>,
34 pub edges: Vec<Edge>,
35 pub node_count: usize,
36}
37
38impl Graph {
39 pub fn from(edges: Vec<Edge>) -> Graph {
40 let mut edge_map = HashMap::new();
41 let mut node_map: HashMap<usize, Vec<Edge>> = HashMap::new();
42 let mut node_count: usize = 0;
43
44 for edge in &edges {
45 edge_map.insert(edge.index.clone(), edge.clone());
46 add_edge_to_node_map(edge.source.clone(), edge.destination.clone(), edge.clone(), &mut node_map);
47 }
48
49 let mut nodes: HashMap<usize, Node> = HashMap::new();
50 for (k, v) in node_map {
51 nodes.insert(k.clone(), Node::from(k.clone(), v.to_vec()));
52 node_count += 1;
53 }
54
55 Graph {
56 nodes_lookup: nodes,
57 edges_lookup: edge_map,
58 node_position_lookup: None,
59 edges,
60 node_count,
61 }
62 }
63
64 pub fn from_adjacency_matrix(matrix: &[&[f32]]) -> Graph {
65 let mut index: usize = 0;
66 let mut vec: Vec<Edge> = Vec::new();
67
68 for (row, array) in matrix.iter().enumerate() {
69 for (col, weight) in array.iter().enumerate() {
70 if !weight.eq(&(0.0 as f32)) {
71 vec.push(Edge::from(index, row, col, weight.clone()));
72 index += 1;
73 }
74 }
75 }
76
77 return Graph::from(vec);
78 }
79
80 pub fn sorted_by_weight_asc(&self) -> Vec<Edge> {
81 let mut sorted_edges = self.edges.clone();
82 sorted_edges.sort_by(|edge1, edge2|
83 edge1.weight.total_cmp(&edge2.weight));
84 return sorted_edges;
85 }
86
87 pub fn offer_positions(&mut self, node_positions: HashMap<usize, Position>) {
88 self.node_position_lookup = Some(node_positions);
89 }
90}
91
92fn add_edge_to_node_map(src: usize, dest: usize, edge: Edge, node_map: &mut HashMap<usize, Vec<Edge>>) {
93 let mut vec = node_map.get(&src).or(Some(&Vec::<Edge>::new())).unwrap().to_vec();
94 vec.push(edge);
95 node_map.insert(src.clone(), vec);
96 node_map.entry(dest).or_insert_with(|| Vec::new());
97}
98
99pub fn minimum_spanning(graph: Graph) -> Graph {
100 let edges = graph.sorted_by_weight_asc();
101 let mut union_find = UnionFind::from(graph.node_count);
102 let mut min_edges = Vec::new();
103
104 for edge in edges {
105 if !union_find.connected(edge.source, edge.destination) {
106 union_find.unify(edge.source, edge.destination);
107 min_edges.push(edge);
108 }
109 }
110
111 return Graph::from(min_edges);
112}
113
114
115#[test]
116fn mst_should_return_graph() {
117 let edge = Edge::from(0, 0, 1, 0.5);
118 let graph = Graph::from(Vec::from([edge]));
119 let min_graph = minimum_spanning(graph);
120
121 assert_eq!(1, min_graph.edges_lookup.keys().count());
122 assert_eq!(2, min_graph.nodes_lookup.keys().count());
123}
124
125#[test]
126fn mst_should_return_graph_with_source_node_having_one_edge() {
127 let edge = Edge::from(0, 0, 1, 0.5);
128 let graph = Graph::from(Vec::from([edge]));
129 let min_graph = minimum_spanning(graph);
130
131 let source_node = min_graph.nodes_lookup.get(&0).unwrap();
132 assert_eq!(1, source_node.edges.to_vec().len());
133 assert!(min_graph.nodes_lookup.contains_key(&0));
134 assert!(min_graph.nodes_lookup.contains_key(&1));
135}
136
137#[test]
138fn mst_should_return_minimum_spanning_tree() {
139 let edge1 = Edge::from(0, 1, 2, 0.0);
140 let edge2 = Edge::from(1, 2, 3, 0.1428571429);
141 let edge3 = Edge::from(2, 1, 0, 0.2857142857);
142 let edge4 = Edge::from(3, 3, 4, 0.2857142857);
143 let edge5 = Edge::from(4, 1, 3, 0.4285714286);
144 let edge6 = Edge::from(5, 0, 3, 0.8571428571);
145 let edge7 = Edge::from(6, 0, 4, 1.0);
146
147
148 let graph = Graph::from(Vec::from([edge1, edge2, edge3, edge4, edge5, edge6, edge7]));
149 let min_graph = minimum_spanning(graph);
150
151 let mut total_cost: f32 = 0.0;
152 for edge in min_graph.edges {
153 total_cost += edge.weight;
154 }
155
156 assert_eq!(0.7142857143, total_cost);
157}
158
159#[test]
160fn edge_from_should_construct_edge() {
161 let edge = Edge::from(0, 2, 3, 0.5);
162
163 assert_eq!(0, edge.index);
164 assert_eq!(2, edge.source);
165 assert_eq!(3, edge.destination);
166 assert_eq!(0.5, edge.weight);
167}
168
169#[test]
170fn sorted_by_weight_asc_should_return_sorted_vec() {
171 let edge3 = Edge::from(2, 2, 3, 0.3);
172 let edge4 = Edge::from(3, 2, 3, 0.7);
173 let edge1 = Edge::from(0, 2, 3, 0.5);
174 let edge2 = Edge::from(1, 2, 3, 0.2);
175
176 let graph = Graph::from(Vec::from([edge1, edge2, edge3, edge4]));
177 let sorted_edges = graph.sorted_by_weight_asc();
178
179 assert_eq!(0.2, sorted_edges[0].weight);
180 assert_eq!(0.3, sorted_edges[1].weight);
181 assert_eq!(0.5, sorted_edges[2].weight);
182 assert_eq!(0.7, sorted_edges[3].weight);
183}
184
185#[test]
186fn create_graph_from_adjacency_matrix() {
187 let matrix: &[&[f32]] = &[
188 &[0.0, 4.0, 0.0, 0.0, 0.0, 0.0, 0.0, 8.0, 0.0],
189 &[4.0, 0.0, 8.0, 0.0, 0.0, 0.0, 0.0, 11.0, 0.0],
190 &[0.0, 8.0, 0.0, 7.0, 0.0, 4.0, 0.0, 0.0, 2.0],
191 &[0.0, 0.0, 7.0, 0.0, 9.0, 14.0, 0.0, 0.0, 0.0],
192 &[0.0, 0.0, 0.0, 9.0, 0.0, 10.0, 0.0, 0.0, 0.0],
193 &[0.0, 0.0, 4.0, 14.0, 10.0, 0.0, 2.0, 0.0, 0.0],
194 &[0.0, 0.0, 0.0, 0.0, 0.0, 2.0, 0.0, 1.0, 6.0],
195 &[8.0, 11.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 7.0],
196 &[0.0, 0.0, 2.0, 0.0, 0.0, 0.0, 6.0, 7.0, 0.0]
197 ];
198
199 let graph = Graph::from_adjacency_matrix(matrix);
200
201 assert_eq!(28, graph.edges.len());
202 assert_eq!(2, graph.nodes_lookup.get(&0).unwrap().edges.len());
203 assert_eq!(3, graph.nodes_lookup.get(&8).unwrap().edges.len());
204 assert_eq!(2.0, graph.nodes_lookup.get(&8).unwrap().edges[0].weight);
205}
206
207#[test]
208fn create_initial_graph_should_not_have_node_positions() {
209 let edge = Edge::from(0, 2, 3, 0.5);
210 let graph = Graph::from(Vec::from([edge]));
211
212 assert!(graph.node_position_lookup.is_none());
213}
214
215#[test]
216fn offer_node_positions_should_set_node_positions() {
217 let edge = Edge::from(0, 2, 3, 0.5);
218 let mut graph = Graph::from(Vec::from([edge.clone()]));
219
220 let mut node_positions: HashMap<usize, Position> = HashMap::new();
221 node_positions.insert((&edge).source.clone(), Position::from(0.3, 0.2, 0.0));
222 node_positions.insert((&edge).destination.clone(), Position::from(0.1, 0.5, 0.0));
223
224 graph.offer_positions(node_positions);
225
226 assert!(graph.node_position_lookup.is_some());
227
228 let position_lookup = graph.node_position_lookup.unwrap();
229 assert_eq!(0.3, position_lookup.get(&(&edge).source).unwrap().x);
230 assert_eq!(0.1, position_lookup.get(&(&edge).destination).unwrap().x);
231}