Skip to content

Commit a119f67

Browse files
committed
ref: refactor mst
1 parent 1af4efa commit a119f67

File tree

1 file changed

+133
-110
lines changed

1 file changed

+133
-110
lines changed

src/graph/minimum_spanning_tree.rs

Lines changed: 133 additions & 110 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,22 @@
1-
use super::DisjointSetUnion;
1+
//! This module implements Kruskal's algorithm to find the Minimum Spanning Tree (MST)
2+
//! of an undirected, weighted graph using a Disjoint Set Union (DSU) for cycle detection.
23
3-
#[derive(Debug)]
4-
pub struct Edge {
5-
source: i64,
6-
destination: i64,
7-
cost: i64,
8-
}
4+
use crate::graph::DisjointSetUnion;
95

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-
}
6+
/// Represents an edge in the graph with a source, destination, and associated cost.
7+
#[derive(Debug, PartialEq, Eq)]
8+
pub struct Edge {
9+
/// The starting vertex of the edge.
10+
source: usize,
11+
/// The ending vertex of the edge.
12+
destination: usize,
13+
/// The cost associated with the edge.
14+
cost: usize,
1615
}
1716

18-
impl Eq for Edge {}
19-
2017
impl Edge {
21-
fn new(source: i64, destination: i64, cost: i64) -> Self {
18+
/// Creates a new edge with the specified source, destination, and cost.
19+
fn new(source: usize, destination: usize, cost: usize) -> Self {
2220
Self {
2321
source,
2422
destination,
@@ -27,112 +25,137 @@ impl Edge {
2725
}
2826
}
2927

30-
pub fn kruskal(mut edges: Vec<Edge>, number_of_vertices: i64) -> (i64, Vec<Edge>) {
31-
let mut dsu = DisjointSetUnion::new(number_of_vertices as usize);
32-
33-
edges.sort_unstable_by(|a, b| a.cost.cmp(&b.cost));
34-
let mut total_cost: i64 = 0;
35-
let mut final_edges: Vec<Edge> = Vec::new();
36-
let mut merge_count: i64 = 0;
37-
for edge in edges.iter() {
38-
if merge_count >= number_of_vertices - 1 {
28+
/// Executes Kruskal's algorithm to compute the Minimum Spanning Tree (MST) of a graph.
29+
///
30+
/// # Parameters
31+
///
32+
/// - `edges`: A vector of `Edge` instances representing all edges in the graph.
33+
/// - `num_vertices`: The total number of vertices in the graph.
34+
///
35+
/// # Returns
36+
///
37+
/// An `Option` containing a tuple with:
38+
///
39+
/// - The total cost of the MST (usize).
40+
/// - A vector of edges that are included in the MST.
41+
///
42+
/// Returns `None` if the graph is disconnected.
43+
///
44+
/// # Complexity
45+
///
46+
/// The time complexity is O(E log E), where E is the number of edges.
47+
pub fn kruskal(mut edges: Vec<Edge>, num_vertices: usize) -> Option<(usize, Vec<Edge>)> {
48+
let mut dsu = DisjointSetUnion::new(num_vertices);
49+
let mut mst_cost: usize = 0;
50+
let mut mst_edges: Vec<Edge> = Vec::with_capacity(num_vertices - 1);
51+
52+
// Sort edges by cost in ascending order
53+
edges.sort_unstable_by_key(|edge| edge.cost);
54+
55+
for edge in edges {
56+
if mst_edges.len() == num_vertices - 1 {
3957
break;
4058
}
4159

42-
let source: i64 = edge.source;
43-
let destination: i64 = edge.destination;
44-
if dsu.merge(source as usize, destination as usize) < usize::MAX {
45-
merge_count += 1;
46-
let cost: i64 = edge.cost;
47-
total_cost += cost;
48-
let final_edge: Edge = Edge::new(source, destination, cost);
49-
final_edges.push(final_edge);
60+
// Attempt to merge the sets containing the edge’s vertices
61+
if dsu.merge(edge.source, edge.destination) != usize::MAX {
62+
mst_cost += edge.cost;
63+
mst_edges.push(edge);
5064
}
5165
}
52-
(total_cost, final_edges)
66+
67+
// Return MST if it includes exactly num_vertices - 1 edges, otherwise None for disconnected graphs
68+
(mst_edges.len() == num_vertices - 1).then_some((mst_cost, mst_edges))
5369
}
5470

5571
#[cfg(test)]
5672
mod tests {
5773
use super::*;
5874

59-
#[test]
60-
fn test_seven_vertices_eleven_edges() {
61-
let edges = vec![
62-
Edge::new(0, 1, 7),
63-
Edge::new(0, 3, 5),
64-
Edge::new(1, 2, 8),
65-
Edge::new(1, 3, 9),
66-
Edge::new(1, 4, 7),
67-
Edge::new(2, 4, 5),
68-
Edge::new(3, 4, 15),
69-
Edge::new(3, 5, 6),
70-
Edge::new(4, 5, 8),
71-
Edge::new(4, 6, 9),
72-
Edge::new(5, 6, 11),
73-
];
74-
75-
let number_of_vertices: i64 = 7;
76-
77-
let expected_total_cost = 39;
78-
let expected_used_edges = vec![
79-
Edge::new(0, 3, 5),
80-
Edge::new(2, 4, 5),
81-
Edge::new(3, 5, 6),
82-
Edge::new(0, 1, 7),
83-
Edge::new(1, 4, 7),
84-
Edge::new(4, 6, 9),
85-
];
86-
87-
let (actual_total_cost, actual_final_edges) = kruskal(edges, number_of_vertices);
88-
89-
assert_eq!(actual_total_cost, expected_total_cost);
90-
assert_eq!(actual_final_edges, expected_used_edges);
75+
macro_rules! test_cases {
76+
($($name:ident: $test_case:expr,)*) => {
77+
$(
78+
#[test]
79+
fn $name() {
80+
let (edges, num_vertices, expected_result) = $test_case;
81+
let actual_result = kruskal(edges, num_vertices);
82+
assert_eq!(actual_result, expected_result);
83+
}
84+
)*
85+
};
9186
}
9287

93-
#[test]
94-
fn test_ten_vertices_twenty_edges() {
95-
let edges = vec![
96-
Edge::new(0, 1, 3),
97-
Edge::new(0, 3, 6),
98-
Edge::new(0, 4, 9),
99-
Edge::new(1, 2, 2),
100-
Edge::new(1, 3, 4),
101-
Edge::new(1, 4, 9),
102-
Edge::new(2, 3, 2),
103-
Edge::new(2, 5, 8),
104-
Edge::new(2, 6, 9),
105-
Edge::new(3, 6, 9),
106-
Edge::new(4, 5, 8),
107-
Edge::new(4, 9, 18),
108-
Edge::new(5, 6, 7),
109-
Edge::new(5, 8, 9),
110-
Edge::new(5, 9, 10),
111-
Edge::new(6, 7, 4),
112-
Edge::new(6, 8, 5),
113-
Edge::new(7, 8, 1),
114-
Edge::new(7, 9, 4),
115-
Edge::new(8, 9, 3),
116-
];
117-
118-
let number_of_vertices: i64 = 10;
119-
120-
let expected_total_cost = 38;
121-
let expected_used_edges = vec![
122-
Edge::new(7, 8, 1),
123-
Edge::new(1, 2, 2),
124-
Edge::new(2, 3, 2),
125-
Edge::new(0, 1, 3),
126-
Edge::new(8, 9, 3),
127-
Edge::new(6, 7, 4),
128-
Edge::new(5, 6, 7),
129-
Edge::new(2, 5, 8),
130-
Edge::new(4, 5, 8),
131-
];
132-
133-
let (actual_total_cost, actual_final_edges) = kruskal(edges, number_of_vertices);
134-
135-
assert_eq!(actual_total_cost, expected_total_cost);
136-
assert_eq!(actual_final_edges, expected_used_edges);
88+
test_cases! {
89+
test_seven_vertices_eleven_edges: (
90+
vec![
91+
Edge::new(0, 1, 7),
92+
Edge::new(0, 3, 5),
93+
Edge::new(1, 2, 8),
94+
Edge::new(1, 3, 9),
95+
Edge::new(1, 4, 7),
96+
Edge::new(2, 4, 5),
97+
Edge::new(3, 4, 15),
98+
Edge::new(3, 5, 6),
99+
Edge::new(4, 5, 8),
100+
Edge::new(4, 6, 9),
101+
Edge::new(5, 6, 11),
102+
],
103+
7,
104+
Some((39, vec![
105+
Edge::new(0, 3, 5),
106+
Edge::new(2, 4, 5),
107+
Edge::new(3, 5, 6),
108+
Edge::new(0, 1, 7),
109+
Edge::new(1, 4, 7),
110+
Edge::new(4, 6, 9),
111+
]))
112+
),
113+
test_ten_vertices_twenty_edges: (
114+
vec![
115+
Edge::new(0, 1, 3),
116+
Edge::new(0, 3, 6),
117+
Edge::new(0, 4, 9),
118+
Edge::new(1, 2, 2),
119+
Edge::new(1, 3, 4),
120+
Edge::new(1, 4, 9),
121+
Edge::new(2, 3, 2),
122+
Edge::new(2, 5, 8),
123+
Edge::new(2, 6, 9),
124+
Edge::new(3, 6, 9),
125+
Edge::new(4, 5, 8),
126+
Edge::new(4, 9, 18),
127+
Edge::new(5, 6, 7),
128+
Edge::new(5, 8, 9),
129+
Edge::new(5, 9, 10),
130+
Edge::new(6, 7, 4),
131+
Edge::new(6, 8, 5),
132+
Edge::new(7, 8, 1),
133+
Edge::new(7, 9, 4),
134+
Edge::new(8, 9, 3),
135+
],
136+
10,
137+
Some((38, vec![
138+
Edge::new(7, 8, 1),
139+
Edge::new(1, 2, 2),
140+
Edge::new(2, 3, 2),
141+
Edge::new(0, 1, 3),
142+
Edge::new(8, 9, 3),
143+
Edge::new(6, 7, 4),
144+
Edge::new(5, 6, 7),
145+
Edge::new(2, 5, 8),
146+
Edge::new(4, 5, 8),
147+
]))
148+
),
149+
test_disconnected_graph: (
150+
vec![
151+
Edge::new(0, 1, 4),
152+
Edge::new(0, 2, 6),
153+
// Component 1 (vertices 0, 1, 2)
154+
Edge::new(3, 4, 2),
155+
// Component 2 (vertices 3, 4)
156+
],
157+
5,
158+
None
159+
),
137160
}
138161
}

0 commit comments

Comments
 (0)