open_hypergraphs/array/vec/
connected_components.rs

1use std::collections::HashMap;
2
3struct UnionFind {
4    /// arbitrarily chosen ancestor of each node
5    parent: Vec<usize>,
6    /// used to order subtrees by depth
7    rank: Vec<usize>,
8}
9
10impl UnionFind {
11    /// Create a new `UnionFind` with n elements, each in their own set.
12    fn new(n: usize) -> Self {
13        Self {
14            parent: (0..n).collect(),
15            rank: vec![0; n],
16        }
17    }
18
19    /// Find the representative (root) of the set containing x.
20    /// Uses path compression to flatten the structure for efficiency.
21    fn find(&mut self, x: usize) -> usize {
22        if self.parent[x] != x {
23            self.parent[x] = self.find(self.parent[x]);
24        }
25        self.parent[x]
26    }
27
28    /// Union the sets containing x and y.
29    /// Uses union by rank to keep the structure shallow.
30    fn union(&mut self, x: usize, y: usize) {
31        let root_x = self.find(x);
32        let root_y = self.find(y);
33        if root_x != root_y {
34            #[allow(clippy::comparison_chain)]
35            if self.rank[root_x] > self.rank[root_y] {
36                self.parent[root_y] = root_x;
37            } else if self.rank[root_x] < self.rank[root_y] {
38                self.parent[root_x] = root_y;
39            } else {
40                self.parent[root_y] = root_x;
41                self.rank[root_x] += 1;
42            }
43        }
44    }
45}
46
47/// Find connected components of a graph stored as a pair of arrays of nodes,
48/// interpreted as a list of edges `sources[i] → targets[i]`
49///
50/// # Panics
51///
52/// * `sources.len() != targets.len()`
53/// * When any `sources[i] >= n` or `targets[i] >= n`
54pub fn connected_components(sources: &[usize], targets: &[usize], n: usize) -> (Vec<usize>, usize) {
55    // Must have equal sized arrays
56    assert_eq!(sources.len(), targets.len());
57    // Arrays must be empty if graph has no nodes.
58    assert!(n > 0 || sources.is_empty());
59    let mut uf = UnionFind::new(n);
60
61    // Union each pair of nodes connected by an edge.
62    for (u, v) in sources.iter().zip(targets) {
63        uf.union(*u, *v);
64    }
65
66    // Find the connected component representative for each node.
67    let node_to_other_node = (0..n).map(|i| uf.find(i)).collect::<Vec<_>>();
68    let (node_to_component_number, num_components) = to_dense(&node_to_other_node);
69    (node_to_component_number, num_components)
70}
71
72#[must_use]
73/// compress a sparse mapping into a set `[0..A)` into a surjective
74/// mapping into a set `[0..num_components)`
75pub fn to_dense(sparse: &[usize]) -> (Vec<usize>, usize) {
76    let mut representative_nodes: HashMap<usize, usize> = HashMap::new();
77    let mut num_components = 0;
78    let mut dense: Vec<usize> = Vec::with_capacity(sparse.len());
79    for a in sparse {
80        if let Some(found_component_number) = representative_nodes.get(a) {
81            dense.push(*found_component_number);
82        } else {
83            representative_nodes.insert(*a, num_components);
84            dense.push(num_components);
85            num_components += 1;
86        }
87    }
88    (dense, num_components)
89}
90
91#[cfg(test)]
92mod test {
93    use proptest::prelude::{Just, Strategy};
94    use proptest::{prop_assert_eq, proptest};
95
96    use super::{connected_components, to_dense};
97
98    #[test]
99    fn to_dense_example() {
100        assert_eq!(to_dense(&[0, 2, 5, 5, 7]), (vec![0, 1, 2, 2, 3], 4));
101    }
102
103    #[test]
104    fn small_graph_components() {
105        let components = connected_components(&[0, 1, 3], &[1, 2, 4], 5);
106        let expected = (vec![0, 0, 0, 1, 1], 2);
107
108        assert_eq!(components, expected);
109    }
110
111    fn edges_strategy() -> impl Strategy<Value = (Vec<usize>, Vec<usize>, usize)> {
112        (1_usize..100).prop_flat_map(|n| {
113            (Just(n), (0..(n + 2) / 2)).prop_flat_map(|(n, e)| {
114                (
115                    proptest::collection::vec(0..n, e),
116                    proptest::collection::vec(0..n, e),
117                    Just(n),
118                )
119            })
120        })
121    }
122
123    proptest! {
124        #[test]
125        fn general_components((sources,targets,num_nodes) in edges_strategy()) {
126            let (mut z, component_count) = connected_components(&sources, &targets, num_nodes);
127            prop_assert_eq!(z.len(),num_nodes);
128            // if there is an edge between two nodes they must be in the same component
129            for (a,b) in sources.into_iter().zip(targets) {
130                prop_assert_eq!(z[a],z[b]);
131            }
132            // there should be the expected number of components
133            let num_components = z.iter().max().copied().expect("At least 1 node") + 1;
134            prop_assert_eq!(num_components,component_count);
135            // the component labels seen should be surjective onto `[0, num_components)`
136            z.sort_unstable();
137            z.dedup();
138            prop_assert_eq!(z, (0..num_components).collect::<Vec<_>>());
139        }
140    }
141}
open_hypergraphs/array/vec/connected_components.rs

open_hypergraphs/array/vec/
connected_components.rs
