Update DSU to use both path compression and union by size (rust-lang#278)

Author: erfan-khadem

src/graph/minimum_spanning_tree.rs

@@ -7,6 +7,12 @@ pub struct Edge {
     cost: i64,
 }
 
+#[derive(Debug)]
+struct DSUNode {
+    parent: i64,
+    subtree_size: i64,
+}
+
 impl PartialEq for Edge {
     fn eq(&self, other: &Self) -> bool {
         self.source == other.source
@@ -27,34 +33,47 @@ impl Edge {
     }
 }
 
-fn make_sets(number_of_vertices: i64) -> Vec<i64> {
-    let mut parent: Vec<i64> = Vec::with_capacity(number_of_vertices as usize);
+fn make_sets(number_of_vertices: i64) -> Vec<DSUNode> {
+    let mut dsu_nodes: Vec<DSUNode> = Vec::with_capacity(number_of_vertices as usize);
     for i in 0..number_of_vertices {
-        parent.push(i);
+        dsu_nodes.push(DSUNode {
+            parent: i,
+            subtree_size: 1,
+        });
     }
-    parent
+    dsu_nodes
 }
 
-fn find(parent: &mut Vec<i64>, x: i64) -> i64 {
+fn find(dsu_nodes: &mut Vec<DSUNode>, x: i64) -> i64 {
     let idx: usize = x as usize;
-    if parent[idx] != x {
-        parent[idx] = find(parent, parent[idx]);
+    if dsu_nodes[idx].parent != x {
+        dsu_nodes[idx].parent = find(dsu_nodes, dsu_nodes[idx].parent);
+        // subtree_size of this vertex might become invalid, but only size of
+        // roots are important and used, so it doesn't matter
     }
-    parent[idx]
+    dsu_nodes[idx].parent
 }
 
-fn merge(parent: &mut Vec<i64>, x: i64, y: i64) {
-    let idx_x: usize = find(parent, x) as usize;
-    let parent_y: i64 = find(parent, y);
-    parent[idx_x] = parent_y;
+fn merge(dsu_nodes: &mut Vec<DSUNode>, x: i64, y: i64) {
+    let mut idx_x: usize = find(dsu_nodes, x) as usize;
+    let mut idx_y: usize = find(dsu_nodes, y) as usize;
+
+    // We should make the smaller root a child of the other
+    // We assume idx_x is the bigger one, and swap it if it is not
+    if dsu_nodes[idx_y].subtree_size > dsu_nodes[idx_x].subtree_size {
+        std::mem::swap(&mut idx_y, &mut idx_x);
+    }
+
+    dsu_nodes[idx_y].parent = idx_x as i64;
+    dsu_nodes[idx_x].subtree_size += dsu_nodes[idx_y].subtree_size;
 }
 
-fn is_same_set(parent: &mut Vec<i64>, x: i64, y: i64) -> bool {
-    find(parent, x) == find(parent, y)
+fn is_same_set(dsu_nodes: &mut Vec<DSUNode>, x: i64, y: i64) -> bool {
+    find(dsu_nodes, x) == find(dsu_nodes, y)
 }
 
 pub fn kruskal(mut edges: Vec<Edge>, number_of_vertices: i64) -> (i64, Vec<Edge>) {
-    let mut parent: Vec<i64> = make_sets(number_of_vertices);
+    let mut dsu_nodes: Vec<DSUNode> = make_sets(number_of_vertices);
 
     edges.sort_unstable_by(|a, b| a.cost.cmp(&b.cost));
     let mut total_cost: i64 = 0;
@@ -67,8 +86,8 @@ pub fn kruskal(mut edges: Vec<Edge>, number_of_vertices: i64) -> (i64, Vec<Edge>
 
         let source: i64 = edge.source;
         let destination: i64 = edge.destination;
-        if !is_same_set(&mut parent, source, destination) {
-            merge(&mut parent, source, destination);
+        if !is_same_set(&mut dsu_nodes, source, destination) {
+            merge(&mut dsu_nodes, source, destination);
             merge_count += 1;
             let cost: i64 = edge.cost;
             total_cost += cost;