Prim算法(最小生成树)

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. fn 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. }

12. // selects a start and run the algorithm from it
13. pub 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. }

22. // only works for a connected graph
23. // if the given graph is not connected it will return the MST of the connected subgraph
24. pub 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();

34. mst.insert(start, BTreeMap::new());

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. }

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. }

47. // the destination is a new vertex
48. add_edge(&mut mst, prev, *t, dist);

50. for (v, c) in &graph[t] {
51. if !mst.contains_key(v) {
52. prio.push(Reverse((*c, v, *t)));
53. }
54. }
55. }

57. mst
58. }

60. #[cfg(test)]
61. mod tests {
62. use super::{add_edge, prim, Graph};
63. use std::collections::BTreeMap;

65. #[test]
66. fn empty() {
67. assert_eq!(prim::<usize, usize>(&BTreeMap::new()), BTreeMap::new());
68. }

70. #[test]
71. fn single_vertex() {
72. let mut graph: Graph<usize, usize> = BTreeMap::new();
73. graph.insert(42, BTreeMap::new());

75. assert_eq!(prim(&graph), graph);
76. }

78. #[test]
79. fn single_edge() {
80. let mut graph = BTreeMap::new();

82. add_edge(&mut graph, 42, 666, 12);

84. assert_eq!(prim(&graph), graph);
85. }

87. #[test]
88. fn tree_1() {
89. let mut graph = BTreeMap::new();

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);

99. assert_eq!(prim(&graph), graph);
100. }

102. #[test]
103. fn tree_2() {
104. let mut graph = BTreeMap::new();

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);

113. assert_eq!(prim(&graph), graph);
114. }

116. #[test]
117. fn tree_3() {
118. let mut graph = BTreeMap::new();

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. }

125. assert_eq!(prim(&graph), graph);
126. }

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);

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);

148. assert_eq!(prim(&graph), ans);
149. }

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);

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);

172. assert_eq!(prim(&graph), ans);
173. }

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);

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);

197. assert_eq!(prim(&graph), ans);
198. }
199. }