最短路径-Dijkstra

1. use std::cmp::Reverse;
2. use std::collections::{BTreeMap, BinaryHeap};
3. use std::ops::Add;

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

7. // performs Dijsktra's algorithm on the given graph from the given start
8. // the graph is a positively-weighted undirected graph
9. //
10. // returns a map that for each reachable vertex associates the distance and the predecessor
11. // since the start has no predecessor but is reachable, map[start] will be None
12. pub fn dijkstra<V: Ord + Copy, E: Ord + Copy + Add<Output = E>>(
13. graph: &Graph<V, E>,
14. start: &V,
15. ) -> BTreeMap<V, Option<(V, E)>> {
16. let mut ans = BTreeMap::new();
17. let mut prio = BinaryHeap::new();

19. // start is the special case that doesn't have a predecessor
20. ans.insert(*start, None);

22. for (new, weight) in &graph[start] {
23. ans.insert(*new, Some((*start, *weight)));
24. prio.push(Reverse((*weight, new, start)));
25. }

27. while let Some(Reverse((dist_new, new, prev))) = prio.pop() {
28. match ans[new] {
29. // what we popped is what is in ans, we'll compute it
30. Some((p, d)) if p == *prev && d == dist_new => {}
31. // otherwise it's not interesting
32. _ => continue,
33. }

35. for (next, weight) in &graph[new] {
36. match ans.get(next) {
37. // if ans[next] is a lower dist than the alternative one, we do nothing
38. Some(Some((_, dist_next))) if dist_new + *weight >= *dist_next => {}
39. // if ans[next] is None then next is start and so the distance won't be changed, it won't be added again in prio
40. Some(None) => {}
41. // the new path is shorter, either new was not in ans or it was farther
42. _ => {
43. ans.insert(*next, Some((*new, *weight + dist_new)));
44. prio.push(Reverse((*weight + dist_new, next, new)));
45. }
46. }
47. }
48. }

50. ans
51. }

53. #[cfg(test)]
54. mod tests {
55. use super::{dijkstra, Graph};
56. use std::collections::BTreeMap;

58. fn add_edge<V: Ord + Copy, E: Ord>(graph: &mut Graph<V, E>, v1: V, v2: V, c: E) {
59. graph.entry(v1).or_insert_with(BTreeMap::new).insert(v2, c);
60. graph.entry(v2).or_insert_with(BTreeMap::new);
61. }

63. #[test]
64. fn single_vertex() {
65. let mut graph: Graph<usize, usize> = BTreeMap::new();
66. graph.insert(0, BTreeMap::new());

68. let mut dists = BTreeMap::new();
69. dists.insert(0, None);

71. assert_eq!(dijkstra(&graph, &0), dists);
72. }

74. #[test]
75. fn single_edge() {
76. let mut graph = BTreeMap::new();
77. add_edge(&mut graph, 0, 1, 2);

79. let mut dists_0 = BTreeMap::new();
80. dists_0.insert(0, None);
81. dists_0.insert(1, Some((0, 2)));

83. assert_eq!(dijkstra(&graph, &0), dists_0);

85. let mut dists_1 = BTreeMap::new();
86. dists_1.insert(1, None);

88. assert_eq!(dijkstra(&graph, &1), dists_1);
89. }

91. #[test]
92. fn tree_1() {
93. let mut graph = BTreeMap::new();
94. let mut dists = BTreeMap::new();
95. dists.insert(1, None);
96. for i in 1..100 {
97. add_edge(&mut graph, i, i * 2, i * 2);
98. add_edge(&mut graph, i, i * 2 + 1, i * 2 + 1);

100. match dists[&i] {
101. Some((_, d)) => {
102. dists.insert(i * 2, Some((i, d + i * 2)));
103. dists.insert(i * 2 + 1, Some((i, d + i * 2 + 1)));
104. }
105. None => {
106. dists.insert(i * 2, Some((i, i * 2)));
107. dists.insert(i * 2 + 1, Some((i, i * 2 + 1)));
108. }
109. }
110. }

112. assert_eq!(dijkstra(&graph, &1), dists);
113. }

115. #[test]
116. fn graph_1() {
117. let mut graph = BTreeMap::new();
118. add_edge(&mut graph, 'a', 'c', 12);
119. add_edge(&mut graph, 'a', 'd', 60);
120. add_edge(&mut graph, 'b', 'a', 10);
121. add_edge(&mut graph, 'c', 'b', 20);
122. add_edge(&mut graph, 'c', 'd', 32);
123. add_edge(&mut graph, 'e', 'a', 7);

125. let mut dists_a = BTreeMap::new();
126. dists_a.insert('a', None);
127. dists_a.insert('c', Some(('a', 12)));
128. dists_a.insert('d', Some(('c', 44)));
129. dists_a.insert('b', Some(('c', 32)));
130. assert_eq!(dijkstra(&graph, &'a'), dists_a);

132. let mut dists_b = BTreeMap::new();
133. dists_b.insert('b', None);
134. dists_b.insert('a', Some(('b', 10)));
135. dists_b.insert('c', Some(('a', 22)));
136. dists_b.insert('d', Some(('c', 54)));
137. assert_eq!(dijkstra(&graph, &'b'), dists_b);

139. let mut dists_c = BTreeMap::new();
140. dists_c.insert('c', None);
141. dists_c.insert('b', Some(('c', 20)));
142. dists_c.insert('d', Some(('c', 32)));
143. dists_c.insert('a', Some(('b', 30)));
144. assert_eq!(dijkstra(&graph, &'c'), dists_c);

146. let mut dists_d = BTreeMap::new();
147. dists_d.insert('d', None);
148. assert_eq!(dijkstra(&graph, &'d'), dists_d);

150. let mut dists_e = BTreeMap::new();
151. dists_e.insert('e', None);
152. dists_e.insert('a', Some(('e', 7)));
153. dists_e.insert('c', Some(('a', 19)));
154. dists_e.insert('d', Some(('c', 51)));
155. dists_e.insert('b', Some(('c', 39)));
156. assert_eq!(dijkstra(&graph, &'e'), dists_e);
157. }
158. }