最短路径-Bellman Ford

1. use std::collections::BTreeMap;
2. use std::ops::Add;

4. use std::ops::Neg;

6. type Graph<V, E> = BTreeMap<V, BTreeMap<V, E>>;

8. // performs the Bellman-Ford algorithm on the given graph from the given start
9. // the graph is an undirected graph
10. //
11. // if there is a negative weighted loop it returns None
12. // else it returns a map that for each reachable vertex associates the distance and the predecessor
13. // since the start has no predecessor but is reachable, map[start] will be None
14. pub fn bellman_ford<
15. V: Ord + Copy,
16. E: Ord + Copy + Add<Output = E> + Neg<Output = E> + std::ops::Sub<Output = E>,
17. >(
18. graph: &Graph<V, E>,
19. start: &V,
20. ) -> Option<BTreeMap<V, Option<(V, E)>>> {
21. let mut ans: BTreeMap<V, Option<(V, E)>> = BTreeMap::new();

23. ans.insert(*start, None);

25. for _ in 1..(graph.len()) {
26. for (u, edges) in graph {
27. let dist_u = match ans.get(u) {
28. Some(Some((_, d))) => Some(*d),
29. Some(None) => None,
30. None => continue,
31. };

33. for (v, d) in edges {
34. match ans.get(v) {
35. Some(Some((_, dist)))
36. // if this is a longer path, do nothing
37. if match dist_u {
38. Some(dist_u) => dist_u + *d >= *dist,
39. None => d >= dist,
40. } => {}
41. Some(None) => {
42. match dist_u {
43. // if dist_u + d < 0 there is a negative loop going by start
44. // else it's just a longer path
45. Some(dist_u) if dist_u >= -*d => {}
46. // negative self edge or negative loop
47. _ => {
48. if *d > *d + *d {
49. return None;
50. }
51. }
52. };
53. }
54. // it's a shorter path: either dist_v was infinite or it was longer than dist_u + d
55. _ => {
56. ans.insert(
57. *v,
58. Some((
59. *u,
60. match dist_u {
61. Some(dist) => dist + *d,
62. None => *d,
63. },
64. )),
65. );
66. }
67. }
68. }
69. }
70. }

72. for (u, edges) in graph {
73. for (v, d) in edges {
74. match (ans.get(u), ans.get(v)) {
75. (Some(None), Some(None)) if *d > *d + *d => return None,
76. (Some(None), Some(Some((_, dv)))) if d < dv => return None,
77. (Some(Some((_, du))), Some(None)) if *du < -*d => return None,
78. (Some(Some((_, du))), Some(Some((_, dv)))) if *du + *d < *dv => return None,
79. (_, _) => {}
80. }
81. }
82. }

84. Some(ans)
85. }

87. #[cfg(test)]
88. mod tests {
89. use super::{bellman_ford, Graph};
90. use std::collections::BTreeMap;

92. fn add_edge<V: Ord + Copy, E: Ord>(graph: &mut Graph<V, E>, v1: V, v2: V, c: E) {
93. graph.entry(v1).or_insert_with(BTreeMap::new).insert(v2, c);
94. graph.entry(v2).or_insert_with(BTreeMap::new);
95. }

97. #[test]
98. fn single_vertex() {
99. let mut graph: Graph<isize, isize> = BTreeMap::new();
100. graph.insert(0, BTreeMap::new());

102. let mut dists = BTreeMap::new();
103. dists.insert(0, None);

105. assert_eq!(bellman_ford(&graph, &0), Some(dists));
106. }

108. #[test]
109. fn single_edge() {
110. let mut graph = BTreeMap::new();
111. add_edge(&mut graph, 0, 1, 2);

113. let mut dists_0 = BTreeMap::new();
114. dists_0.insert(0, None);
115. dists_0.insert(1, Some((0, 2)));

117. assert_eq!(bellman_ford(&graph, &0), Some(dists_0));

119. let mut dists_1 = BTreeMap::new();
120. dists_1.insert(1, None);

122. assert_eq!(bellman_ford(&graph, &1), Some(dists_1));
123. }

125. #[test]
126. fn tree_1() {
127. let mut graph = BTreeMap::new();
128. let mut dists = BTreeMap::new();
129. dists.insert(1, None);
130. for i in 1..100 {
131. add_edge(&mut graph, i, i * 2, i * 2);
132. add_edge(&mut graph, i, i * 2 + 1, i * 2 + 1);

134. match dists[&i] {
135. Some((_, d)) => {
136. dists.insert(i * 2, Some((i, d + i * 2)));
137. dists.insert(i * 2 + 1, Some((i, d + i * 2 + 1)));
138. }
139. None => {
140. dists.insert(i * 2, Some((i, i * 2)));
141. dists.insert(i * 2 + 1, Some((i, i * 2 + 1)));
142. }
143. }
144. }

146. assert_eq!(bellman_ford(&graph, &1), Some(dists));
147. }

149. #[test]
150. fn graph_1() {
151. let mut graph = BTreeMap::new();
152. add_edge(&mut graph, 'a', 'c', 12);
153. add_edge(&mut graph, 'a', 'd', 60);
154. add_edge(&mut graph, 'b', 'a', 10);
155. add_edge(&mut graph, 'c', 'b', 20);
156. add_edge(&mut graph, 'c', 'd', 32);
157. add_edge(&mut graph, 'e', 'a', 7);

159. let mut dists_a = BTreeMap::new();
160. dists_a.insert('a', None);
161. dists_a.insert('c', Some(('a', 12)));
162. dists_a.insert('d', Some(('c', 44)));
163. dists_a.insert('b', Some(('c', 32)));
164. assert_eq!(bellman_ford(&graph, &'a'), Some(dists_a));

166. let mut dists_b = BTreeMap::new();
167. dists_b.insert('b', None);
168. dists_b.insert('a', Some(('b', 10)));
169. dists_b.insert('c', Some(('a', 22)));
170. dists_b.insert('d', Some(('c', 54)));
171. assert_eq!(bellman_ford(&graph, &'b'), Some(dists_b));

173. let mut dists_c = BTreeMap::new();
174. dists_c.insert('c', None);
175. dists_c.insert('b', Some(('c', 20)));
176. dists_c.insert('d', Some(('c', 32)));
177. dists_c.insert('a', Some(('b', 30)));
178. assert_eq!(bellman_ford(&graph, &'c'), Some(dists_c));

180. let mut dists_d = BTreeMap::new();
181. dists_d.insert('d', None);
182. assert_eq!(bellman_ford(&graph, &'d'), Some(dists_d));

184. let mut dists_e = BTreeMap::new();
185. dists_e.insert('e', None);
186. dists_e.insert('a', Some(('e', 7)));
187. dists_e.insert('c', Some(('a', 19)));
188. dists_e.insert('d', Some(('c', 51)));
189. dists_e.insert('b', Some(('c', 39)));
190. assert_eq!(bellman_ford(&graph, &'e'), Some(dists_e));
191. }

193. #[test]
194. fn graph_2() {
195. let mut graph = BTreeMap::new();
196. add_edge(&mut graph, 0, 1, 6);
197. add_edge(&mut graph, 0, 3, 7);
198. add_edge(&mut graph, 1, 2, 5);
199. add_edge(&mut graph, 1, 3, 8);
200. add_edge(&mut graph, 1, 4, -4);
201. add_edge(&mut graph, 2, 1, -2);
202. add_edge(&mut graph, 3, 2, -3);
203. add_edge(&mut graph, 3, 4, 9);
204. add_edge(&mut graph, 4, 0, 3);
205. add_edge(&mut graph, 4, 2, 7);

207. let mut dists_0 = BTreeMap::new();
208. dists_0.insert(0, None);
209. dists_0.insert(1, Some((2, 2)));
210. dists_0.insert(2, Some((3, 4)));
211. dists_0.insert(3, Some((0, 7)));
212. dists_0.insert(4, Some((1, -2)));
213. assert_eq!(bellman_ford(&graph, &0), Some(dists_0));

215. let mut dists_1 = BTreeMap::new();
216. dists_1.insert(0, Some((4, -1)));
217. dists_1.insert(1, None);
218. dists_1.insert(2, Some((4, 3)));
219. dists_1.insert(3, Some((0, 6)));
220. dists_1.insert(4, Some((1, -4)));
221. assert_eq!(bellman_ford(&graph, &1), Some(dists_1));

223. let mut dists_2 = BTreeMap::new();
224. dists_2.insert(0, Some((4, -3)));
225. dists_2.insert(1, Some((2, -2)));
226. dists_2.insert(2, None);
227. dists_2.insert(3, Some((0, 4)));
228. dists_2.insert(4, Some((1, -6)));
229. assert_eq!(bellman_ford(&graph, &2), Some(dists_2));

231. let mut dists_3 = BTreeMap::new();
232. dists_3.insert(0, Some((4, -6)));
233. dists_3.insert(1, Some((2, -5)));
234. dists_3.insert(2, Some((3, -3)));
235. dists_3.insert(3, None);
236. dists_3.insert(4, Some((1, -9)));
237. assert_eq!(bellman_ford(&graph, &3), Some(dists_3));

239. let mut dists_4 = BTreeMap::new();
240. dists_4.insert(0, Some((4, 3)));
241. dists_4.insert(1, Some((2, 5)));
242. dists_4.insert(2, Some((4, 7)));
243. dists_4.insert(3, Some((0, 10)));
244. dists_4.insert(4, None);
245. assert_eq!(bellman_ford(&graph, &4), Some(dists_4));
246. }

248. #[test]
249. fn graph_with_negative_loop() {
250. let mut graph = BTreeMap::new();
251. add_edge(&mut graph, 0, 1, 6);
252. add_edge(&mut graph, 0, 3, 7);
253. add_edge(&mut graph, 1, 2, 5);
254. add_edge(&mut graph, 1, 3, 8);
255. add_edge(&mut graph, 1, 4, -4);
256. add_edge(&mut graph, 2, 1, -4);
257. add_edge(&mut graph, 3, 2, -3);
258. add_edge(&mut graph, 3, 4, 9);
259. add_edge(&mut graph, 4, 0, 3);
260. add_edge(&mut graph, 4, 2, 7);

262. assert_eq!(bellman_ford(&graph, &0), None);
263. assert_eq!(bellman_ford(&graph, &1), None);
264. assert_eq!(bellman_ford(&graph, &2), None);
265. assert_eq!(bellman_ford(&graph, &3), None);
266. assert_eq!(bellman_ford(&graph, &4), None);
267. }
268. }