utils/
disjoint_set.rs

1//! Disjoint-set data structures.
2
3/// Disjoint-set union (DSU) data structure.
4///
5/// # Example
6/// ```
7/// # use utils::disjoint_set::Dsu;
8/// let mut dsu = Dsu::new(5);
9///
10/// for i in 0..5 {
11///     assert_eq!(dsu.find(i), i);
12///     assert_eq!(dsu.set_size(i), 1);
13/// }
14///
15/// assert!(dsu.union(0, 1));
16/// assert!(dsu.union(1, 2));
17/// assert!(!dsu.union(0, 2));
18/// assert!(dsu.union(3, 4));
19///
20/// assert_eq!(dsu.find(0), dsu.find(1));
21/// assert_eq!(dsu.find(0), dsu.find(2));
22/// assert_eq!(dsu.set_size(0), 3);
23///
24/// assert_eq!(dsu.find(3), dsu.find(4));
25/// assert_eq!(dsu.set_size(3), 2);
26/// assert_ne!(dsu.find(0), dsu.find(3));
27///
28/// assert!(dsu.union(1, 3));
29/// assert_eq!(dsu.set_size(0), 5);
30/// for i in 1..5 {
31///     assert_eq!(dsu.find(0), dsu.find(i));
32/// }
33/// ```
34#[derive(Debug)]
35pub struct Dsu {
36    parent: Vec<usize>,
37    size: Vec<usize>,
38}
39
40impl Dsu {
41    /// Creates a new disjoint-set union data structure with `n` elements.
42    #[must_use]
43    pub fn new(n: usize) -> Self {
44        Dsu {
45            parent: (0..n).collect(),
46            size: vec![1; n],
47        }
48    }
49
50    /// Find the root of the set containing the element at index `x`.
51    ///
52    /// This requires `&mut self` because it performs path compression, making later calls faster.
53    #[inline]
54    #[must_use]
55    pub fn find(&mut self, mut x: usize) -> usize {
56        let mut root = x;
57        while self.parent[root] != root {
58            root = self.parent[root];
59        }
60        while self.parent[x] != root {
61            (x, self.parent[x]) = (self.parent[x], root);
62        }
63        root
64    }
65
66    /// Merges the sets containing index `a` and `b`.
67    ///
68    /// Returns `true` if the sets were merged, `false` if they were already in the same set.
69    #[inline]
70    pub fn union(&mut self, a: usize, b: usize) -> bool {
71        let mut ra = self.find(a);
72        let mut rb = self.find(b);
73        if ra == rb {
74            return false;
75        }
76        if self.size[ra] < self.size[rb] {
77            (ra, rb) = (rb, ra);
78        }
79        self.parent[rb] = ra;
80        self.size[ra] += self.size[rb];
81        true
82    }
83
84    /// Returns the size of the set containing the element at index `x`.
85    ///
86    /// This requires `&mut self` because it performs path compression, making later calls faster.
87    #[inline]
88    #[must_use]
89    pub fn set_size(&mut self, x: usize) -> usize {
90        let root = self.find(x);
91        self.size[root]
92    }
93
94    /// Returns the size of the set with root `x`.
95    ///
96    /// This function will panic if `x` is not a root.
97    #[inline]
98    #[must_use]
99    pub fn root_size(&self, x: usize) -> usize {
100        assert_eq!(self.parent[x], x);
101        self.size[x]
102    }
103
104    /// Returns an iterator over the roots of the disjoint-set.
105    #[inline]
106    pub fn roots(&self) -> impl Iterator<Item = usize> + '_ {
107        (0..self.parent.len()).filter(|&x| self.parent[x] == x)
108    }
109}
utils/disjoint_set.rs

utils/
disjoint_set.rs
