arris_math/
disjoint_sets.rs

1//! Disjoint sets over `0..n`: the connected components of a graph the
2//! caller walks, labelled the same way on every platform.
3
4/// Disjoint sets over `0..n`, the lowest member of a set its root, so a
5/// numbering derived from the roots depends on the elements' order and on
6/// nothing else.
7///
8/// ```
9/// use arris_math::DisjointSets;
10///
11/// let mut sets = DisjointSets::new(5);
12/// assert!(sets.union(3, 1));
13/// assert!(sets.union(4, 3));
14/// assert!(!sets.union(1, 4)); // already together
15/// assert_eq!((sets.find(4), sets.find(3), sets.find(2)), (1, 1, 2));
16/// ```
17#[derive(Debug, Clone)]
18pub struct DisjointSets(Vec<usize>);
19
20impl DisjointSets {
21    /// `n` singletons.
22    pub fn new(n: usize) -> Self {
23        DisjointSets((0..n).collect())
24    }
25
26    /// The root of `i`'s set: its lowest member. Halves the path it walks.
27    ///
28    /// Panics if `i` is not below the `n` of [`DisjointSets::new`]: an
29    /// index is the caller's own, never geometry.
30    pub fn find(&mut self, mut i: usize) -> usize {
31        while self.0[i] != i {
32            self.0[i] = self.0[self.0[i]];
33            i = self.0[i];
34        }
35        i
36    }
37
38    /// Joins the sets of `a` and `b`; `false` if they were one already.
39    pub fn union(&mut self, a: usize, b: usize) -> bool {
40        let (a, b) = (self.find(a), self.find(b));
41        if a == b {
42            return false;
43        }
44        self.0[a.max(b)] = a.min(b);
45        true
46    }
47}
48
49#[cfg(test)]
50mod tests {
51    use super::*;
52    use proptest::prelude::*;
53
54    proptest! {
55        /// Against a naive relabelling partition: the same sets, each
56        /// rooted at its least member.
57        #[test]
58        fn matches_a_naive_partition(
59            n in 1usize..40,
60            pairs in prop::collection::vec((0usize..40, 0usize..40), 0..60),
61        ) {
62            let mut sets = DisjointSets::new(n);
63            let mut label: Vec<usize> = (0..n).collect();
64            for (a, b) in pairs {
65                let (a, b) = (a % n, b % n);
66                let merged = sets.union(a, b);
67                let (la, lb) = (label[a], label[b]);
68                prop_assert_eq!(merged, la != lb);
69                let (keep, drop) = (la.min(lb), la.max(lb));
70                for l in &mut label {
71                    if *l == drop {
72                        *l = keep;
73                    }
74                }
75            }
76            for (i, &l) in label.iter().enumerate() {
77                prop_assert_eq!(sets.find(i), l);
78            }
79        }
80    }
81}
arris_math/disjoint_sets.rs

arris_math/
disjoint_sets.rs
