ac_library/
dsu.rs

1//! A Disjoint set union (DSU) with union by size and path compression.
2
3/// A Disjoint set union (DSU) with union by size and path compression.
4///
5/// See: [Zvi Galil and Giuseppe F. Italiano, Data structures and algorithms for disjoint set union problems](https://core.ac.uk/download/pdf/161439519.pdf)
6///
7/// In the following documentation, let $n$ be the size of the DSU.
8///
9/// # Example
10///
11/// ```
12/// use ac_library::Dsu;
13/// use proconio::{input, source::once::OnceSource};
14///
15/// input! {
16///     from OnceSource::from(
17///         "5\n\
18///          3\n\
19///          0 1\n\
20///          2 3\n\
21///          3 4\n",
22///     ),
23///     n: usize,
24///     abs: [(usize, usize)],
25/// }
26///
27/// let mut dsu = Dsu::new(n);
28/// for (a, b) in abs {
29///     dsu.merge(a, b);
30/// }
31///
32/// assert!(dsu.same(0, 1));
33/// assert!(!dsu.same(1, 2));
34/// assert!(dsu.same(2, 4));
35///
36/// assert_eq!(
37///     dsu.groups()
38///         .into_iter()
39///         .map(|mut group| {
40///             group.sort_unstable();
41///             group
42///         })
43///         .collect::<Vec<_>>(),
44///     [&[0, 1][..], &[2, 3, 4][..]],
45/// );
46/// ```
47#[derive(Clone, Debug)]
48pub struct Dsu {
49    n: usize,
50    // root node: -1 * component size
51    // otherwise: parent
52    parent_or_size: Vec<i32>,
53}
54
55impl Dsu {
56    /// Creates a new `Dsu`.
57    ///
58    /// # Constraints
59    ///
60    /// - $0 \leq n \leq 10^8$
61    ///
62    /// # Complexity
63    ///
64    /// - $O(n)$
65    pub fn new(size: usize) -> Self {
66        Self {
67            n: size,
68            parent_or_size: vec![-1; size],
69        }
70    }
71
72    // `\textsc` does not work in KaTeX
73    /// Performs the Uɴɪᴏɴ operation.
74    ///
75    /// # Constraints
76    ///
77    /// - $0 \leq a < n$
78    /// - $0 \leq b < n$
79    ///
80    /// # Panics
81    ///
82    /// Panics if the above constraints are not satisfied.
83    ///
84    /// # Complexity
85    ///
86    /// - $O(\alpha(n))$ amortized
87    pub fn merge(&mut self, a: usize, b: usize) -> usize {
88        assert!(a < self.n);
89        assert!(b < self.n);
90        let (mut x, mut y) = (self.leader(a), self.leader(b));
91        if x == y {
92            return x;
93        }
94        if -self.parent_or_size[x] < -self.parent_or_size[y] {
95            std::mem::swap(&mut x, &mut y);
96        }
97        self.parent_or_size[x] += self.parent_or_size[y];
98        self.parent_or_size[y] = x as i32;
99        x
100    }
101
102    /// Returns whether the vertices $a$ and $b$ are in the same connected component.
103    ///
104    /// # Constraints
105    ///
106    /// - $0 \leq a < n$
107    /// - $0 \leq b < n$
108    ///
109    /// # Panics
110    ///
111    /// Panics if the above constraint is not satisfied.
112    ///
113    /// # Complexity
114    ///
115    /// - $O(\alpha(n))$ amortized
116    pub fn same(&mut self, a: usize, b: usize) -> bool {
117        assert!(a < self.n);
118        assert!(b < self.n);
119        self.leader(a) == self.leader(b)
120    }
121
122    /// Performs the Fɪɴᴅ operation.
123    ///
124    /// # Constraints
125    ///
126    /// - $0 \leq a < n$
127    ///
128    /// # Panics
129    ///
130    /// Panics if the above constraint is not satisfied.
131    ///
132    /// # Complexity
133    ///
134    /// - $O(\alpha(n))$ amortized
135    pub fn leader(&mut self, a: usize) -> usize {
136        assert!(a < self.n);
137        self._leader(a)
138    }
139
140    /// Returns the size of the connected component that contains the vertex $a$.
141    ///
142    /// # Constraints
143    ///
144    /// - $0 \leq a < n$
145    ///
146    /// # Panics
147    ///
148    /// Panics if the above constraint is not satisfied.
149    ///
150    /// # Complexity
151    ///
152    /// - $O(\alpha(n))$ amortized
153    pub fn size(&mut self, a: usize) -> usize {
154        assert!(a < self.n);
155        let x = self.leader(a);
156        -self.parent_or_size[x] as usize
157    }
158
159    /// Divides the graph into connected components.
160    ///
161    /// The result may not be ordered.
162    ///
163    /// # Complexity
164    ///
165    /// - $O(n)$
166    pub fn groups(&mut self) -> Vec<Vec<usize>> {
167        let mut leader_buf = vec![0; self.n];
168        let mut group_size = vec![0; self.n];
169        for i in 0..self.n {
170            leader_buf[i] = self.leader(i);
171            group_size[leader_buf[i]] += 1;
172        }
173        let mut result = vec![Vec::new(); self.n];
174        for i in 0..self.n {
175            result[i].reserve(group_size[i]);
176        }
177        for i in 0..self.n {
178            result[leader_buf[i]].push(i);
179        }
180        result
181            .into_iter()
182            .filter(|x| !x.is_empty())
183            .collect::<Vec<Vec<usize>>>()
184    }
185
186    fn _leader(&mut self, a: usize) -> usize {
187        if self.parent_or_size[a] < 0 {
188            return a;
189        }
190        self.parent_or_size[a] = self._leader(self.parent_or_size[a] as usize) as i32;
191        self.parent_or_size[a] as usize
192    }
193}
194
195#[cfg(test)]
196mod tests {
197    use super::*;
198
199    #[test]
200    fn dsu_works() {
201        let mut d = Dsu::new(4);
202        d.merge(0, 1);
203        assert!(d.same(0, 1));
204        d.merge(1, 2);
205        assert!(d.same(0, 2));
206        assert_eq!(d.size(0), 3);
207        assert!(!d.same(0, 3));
208        assert_eq!(d.groups(), vec![vec![0, 1, 2], vec![3]]);
209    }
210}
ac_library/dsu.rs

ac_library/
dsu.rs
