最小生成树

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