最小生成树

1. use std::vec::Vec;

3. #[derive(Debug)]
4. pub struct Edge {
5. source: i64,
6. destination: i64,
7. cost: i64,
8. }

10. impl PartialEq for Edge {
11. fn eq(&self, other: &Self) -> bool {
12. self.source == other.source
13. && self.destination == other.destination
14. && self.cost == other.cost
15. }
16. }

18. impl Eq for Edge {}

20. impl Edge {
21. fn new(source: i64, destination: i64, cost: i64) -> Self {
22. Self {
23. source,
24. destination,
25. cost,
26. }
27. }
28. }

30. fn make_sets(number_of_vertices: i64) -> Vec<i64> {
31. let mut parent: Vec<i64> = Vec::with_capacity(number_of_vertices as usize);
32. for i in 0..number_of_vertices {
33. parent.push(i);
34. }
35. parent
36. }

38. fn find(parent: &mut Vec<i64>, x: i64) -> i64 {
39. let idx: usize = x as usize;
40. if parent[idx] != x {
41. parent[idx] = find(parent, parent[idx]);
42. }
43. parent[idx]
44. }

46. fn merge(parent: &mut Vec<i64>, x: i64, y: i64) {
47. let idx_x: usize = find(parent, x) as usize;
48. let parent_y: i64 = find(parent, y);
49. parent[idx_x] = parent_y;
50. }

52. fn is_same_set(parent: &mut Vec<i64>, x: i64, y: i64) -> bool {
53. find(parent, x) == find(parent, y)
54. }

56. pub fn kruskal(mut edges: Vec<Edge>, number_of_vertices: i64) -> (i64, Vec<Edge>) {
57. let mut parent: Vec<i64> = make_sets(number_of_vertices);

59. edges.sort_unstable_by(|a, b| a.cost.cmp(&b.cost));
60. let mut total_cost: i64 = 0;
61. let mut final_edges: Vec<Edge> = Vec::new();
62. let mut merge_count: i64 = 0;
63. for edge in edges.iter() {
64. if merge_count >= number_of_vertices - 1 {
65. break;
66. }

68. let source: i64 = edge.source;
69. let destination: i64 = edge.destination;
70. if !is_same_set(&mut parent, source, destination) {
71. merge(&mut parent, source, destination);
72. merge_count += 1;
73. let cost: i64 = edge.cost;
74. total_cost += cost;
75. let final_edge: Edge = Edge::new(source, destination, cost);
76. final_edges.push(final_edge);
77. }
78. }
79. (total_cost, final_edges)
80. }

82. #[cfg(test)]
83. mod tests {
84. use super::*;

86. #[test]
87. fn test_seven_vertices_eleven_edges() {
88. let mut edges: Vec<Edge> = Vec::new();
89. edges.push(Edge::new(0, 1, 7));
90. edges.push(Edge::new(0, 3, 5));
91. edges.push(Edge::new(1, 2, 8));
92. edges.push(Edge::new(1, 3, 9));
93. edges.push(Edge::new(1, 4, 7));
94. edges.push(Edge::new(2, 4, 5));
95. edges.push(Edge::new(3, 4, 15));
96. edges.push(Edge::new(3, 5, 6));
97. edges.push(Edge::new(4, 5, 8));
98. edges.push(Edge::new(4, 6, 9));
99. edges.push(Edge::new(5, 6, 11));

101. let number_of_vertices: i64 = 7;

103. let expected_total_cost = 39;
104. let mut expected_used_edges: Vec<Edge> = Vec::new();
105. expected_used_edges.push(Edge::new(0, 3, 5));
106. expected_used_edges.push(Edge::new(2, 4, 5));
107. expected_used_edges.push(Edge::new(3, 5, 6));
108. expected_used_edges.push(Edge::new(0, 1, 7));
109. expected_used_edges.push(Edge::new(1, 4, 7));
110. expected_used_edges.push(Edge::new(4, 6, 9));

112. let (actual_total_cost, actual_final_edges) = kruskal(edges, number_of_vertices);

114. assert_eq!(actual_total_cost, expected_total_cost);
115. assert_eq!(actual_final_edges, expected_used_edges);
116. }

118. #[test]
119. fn test_ten_vertices_twenty_edges() {
120. let mut edges: Vec<Edge> = Vec::new();
121. edges.push(Edge::new(0, 1, 3));
122. edges.push(Edge::new(0, 3, 6));
123. edges.push(Edge::new(0, 4, 9));
124. edges.push(Edge::new(1, 2, 2));
125. edges.push(Edge::new(1, 3, 4));
126. edges.push(Edge::new(1, 4, 9));
127. edges.push(Edge::new(2, 3, 2));
128. edges.push(Edge::new(2, 5, 8));
129. edges.push(Edge::new(2, 6, 9));
130. edges.push(Edge::new(3, 6, 9));
131. edges.push(Edge::new(4, 5, 8));
132. edges.push(Edge::new(4, 9, 18));
133. edges.push(Edge::new(5, 6, 7));
134. edges.push(Edge::new(5, 8, 9));
135. edges.push(Edge::new(5, 9, 10));
136. edges.push(Edge::new(6, 7, 4));
137. edges.push(Edge::new(6, 8, 5));
138. edges.push(Edge::new(7, 8, 1));
139. edges.push(Edge::new(7, 9, 4));
140. edges.push(Edge::new(8, 9, 3));

142. let number_of_vertices: i64 = 10;

144. let expected_total_cost = 38;
145. let mut expected_used_edges = Vec::new();
146. expected_used_edges.push(Edge::new(7, 8, 1));
147. expected_used_edges.push(Edge::new(1, 2, 2));
148. expected_used_edges.push(Edge::new(2, 3, 2));
149. expected_used_edges.push(Edge::new(0, 1, 3));
150. expected_used_edges.push(Edge::new(8, 9, 3));
151. expected_used_edges.push(Edge::new(6, 7, 4));
152. expected_used_edges.push(Edge::new(5, 6, 7));
153. expected_used_edges.push(Edge::new(2, 5, 8));
154. expected_used_edges.push(Edge::new(4, 5, 8));

156. let (actual_total_cost, actual_final_edges) = kruskal(edges, number_of_vertices);

158. assert_eq!(actual_total_cost, expected_total_cost);
159. assert_eq!(actual_final_edges, expected_used_edges);
160. }
161. }