1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
//! Union-find over `0..n` with path compression and union by size.
//!
//! Shared infrastructure: graph minimum spanning trees, percolation cluster
//! labelling, and single-linkage clustering all reduce to the same
//! "merge these two, are these two together" question.
/// Disjoint-set forest over the elements `0..n`.
#[derive(Debug, Clone)]
pub struct DisjointSet {
/// `parent[i]` is `i` itself for a root, otherwise the next node up.
parent: Vec<usize>,
/// Number of elements in the tree rooted here. Meaningful only at roots.
size: Vec<usize>,
/// Number of disjoint sets currently represented.
count: usize,
}
impl DisjointSet {
/// `n` singleton sets.
#[must_use]
pub fn new(n: usize) -> Self {
Self {
parent: (0..n).collect(),
size: vec![1; n],
count: n,
}
}
/// Number of elements the structure was built over.
#[must_use]
pub fn len(&self) -> usize {
self.parent.len()
}
/// True when built over zero elements.
#[must_use]
pub fn is_empty(&self) -> bool {
self.parent.is_empty()
}
/// Number of disjoint sets.
///
/// Starts at `n` and drops by one on every union that actually merges.
#[must_use]
pub fn count(&self) -> usize {
self.count
}
/// Representative of `x`'s set, compressing the path as it climbs.
///
/// Iterative rather than recursive: a degenerate forest built by
/// `union_unbalanced`-style calls could otherwise overflow the stack, and
/// this is called in inner loops.
pub fn find(&mut self, x: usize) -> usize {
let mut root = x;
while self.parent[root] != root {
root = self.parent[root];
}
// Second pass: point every node on the path straight at the root.
let mut cur = x;
while self.parent[cur] != root {
let next = self.parent[cur];
self.parent[cur] = root;
cur = next;
}
root
}
/// Merges the sets containing `a` and `b`.
///
/// Returns `true` when they were previously separate, so a caller can
/// count merges (Kruskal accepts exactly the edges for which this is
/// true).
pub fn union(&mut self, a: usize, b: usize) -> bool {
let (mut ra, mut rb) = (self.find(a), self.find(b));
if ra == rb {
return false;
}
// Hang the smaller tree under the larger, which bounds the height by
// log2(n) even before path compression.
if self.size[ra] < self.size[rb] {
std::mem::swap(&mut ra, &mut rb);
}
self.parent[rb] = ra;
self.size[ra] += self.size[rb];
self.count -= 1;
true
}
/// True when `a` and `b` lie in the same set.
pub fn connected(&mut self, a: usize, b: usize) -> bool {
self.find(a) == self.find(b)
}
/// Size of the set containing `x`.
pub fn set_size(&mut self, x: usize) -> usize {
let r = self.find(x);
self.size[r]
}
/// The sets, each as a sorted list of members, ordered by first member.
pub fn sets(&mut self) -> Vec<Vec<usize>> {
let n = self.len();
let mut by_root: std::collections::HashMap<usize, Vec<usize>> =
std::collections::HashMap::new();
for i in 0..n {
let r = self.find(i);
by_root.entry(r).or_default().push(i);
}
let mut out: Vec<Vec<usize>> = by_root.into_values().collect();
out.sort_by_key(|s| s[0]);
out
}
/// A labelling in `0..count()` that is constant on each set.
///
/// Labels are assigned in order of each set's smallest member, so the
/// result depends only on the partition and not on the union order.
pub fn labels(&mut self) -> Vec<usize> {
let n = self.len();
let mut label = vec![usize::MAX; n];
let mut next = 0usize;
for i in 0..n {
let r = self.find(i);
if label[r] == usize::MAX {
label[r] = next;
next += 1;
}
label[i] = label[r];
}
label
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::monte_carlo::Rng;
/// A value in `0..n` taken from the high bits.
///
/// `next_u64() % n` would read the low bits of a linear congruential
/// generator, where bit `b` has period `2^(b+1)`; on a small `n` that
/// cycles through a handful of values and would leave most pairs untried.
fn pick(rng: &mut Rng, n: usize) -> usize {
((u128::from(rng.next_u64()) * n as u128) >> 64) as usize
}
/// The structure must agree with the equivalence relation generated by
/// the same unions, computed by transitive closure.
#[test]
fn agrees_with_transitive_closure() {
const N: usize = 40;
let mut rng = Rng::new(0x51D5_u64);
let mut ds = DisjointSet::new(N);
// Reference: a dense reachability matrix closed under composition.
let mut reach = vec![vec![false; N]; N];
for (i, row) in reach.iter_mut().enumerate() {
row[i] = true;
}
for _ in 0..120 {
let a = pick(&mut rng, N);
let b = pick(&mut rng, N);
ds.union(a, b);
// Close the reference by hand: everything reaching a now reaches
// everything b reaches, and vice versa.
let ca: Vec<usize> = (0..N).filter(|&i| reach[i][a]).collect();
let cb: Vec<usize> = (0..N).filter(|&i| reach[i][b]).collect();
for &i in &ca {
for &j in &cb {
reach[i][j] = true;
reach[j][i] = true;
}
}
for i in 0..N {
for j in 0..N {
assert_eq!(
ds.connected(i, j),
reach[i][j],
"disagreement on ({i}, {j})"
);
}
}
}
}
/// count() is exactly the number of sets, and the sets partition 0..n.
#[test]
fn count_and_sets_form_a_partition() {
const N: usize = 50;
let mut rng = Rng::new(7);
let mut ds = DisjointSet::new(N);
assert_eq!(ds.count(), N);
for _ in 0..80 {
let a = pick(&mut rng, N);
let b = pick(&mut rng, N);
ds.union(a, b);
let sets = ds.sets();
assert_eq!(sets.len(), ds.count());
// Disjoint and covering: the sizes sum to N and every element
// appears once.
let total: usize = sets.iter().map(|s| s.len()).sum();
assert_eq!(total, N);
let mut seen = [false; N];
for s in &sets {
for &x in s {
assert!(!seen[x], "{x} appears in two sets");
seen[x] = true;
}
}
// set_size agrees with the enumerated set.
for s in &sets {
for &x in s {
assert_eq!(ds.set_size(x), s.len());
}
}
}
}
/// A union that merges returns true exactly once per merge, so the number
/// of true returns is n - count().
#[test]
fn merges_counted_exactly() {
const N: usize = 30;
let mut rng = Rng::new(99);
let mut ds = DisjointSet::new(N);
let mut merges = 0usize;
for _ in 0..200 {
let a = pick(&mut rng, N);
let b = pick(&mut rng, N);
if ds.union(a, b) {
merges += 1;
}
}
assert_eq!(merges, N - ds.count());
// Everything is joined by 200 random unions on 30 elements with
// overwhelming probability; assert the weaker invariant that holds
// regardless.
assert!(ds.count() >= 1);
assert_eq!(ds.count(), N - merges);
}
/// labels() depends only on the partition, not on the order of unions.
#[test]
fn labels_are_order_independent() {
let mut a = DisjointSet::new(9);
for (x, y) in [(0, 3), (3, 6), (1, 4), (4, 7), (2, 5)] {
a.union(x, y);
}
let mut b = DisjointSet::new(9);
// Same partition, unions applied in a different order and direction.
for (x, y) in [(5, 2), (7, 1), (6, 0), (4, 1), (3, 0)] {
b.union(x, y);
}
assert_eq!(a.labels(), b.labels());
// And the labelling is a surjection onto 0..count.
let labels = a.labels();
let mut distinct: Vec<usize> = labels.clone();
distinct.sort_unstable();
distinct.dedup();
assert_eq!(distinct, (0..a.count()).collect::<Vec<_>>());
}
/// find() must be idempotent and constant across a set: the representative
/// is a function of the set, not of the query.
#[test]
fn representative_is_a_function_of_the_set() {
let mut ds = DisjointSet::new(20);
for i in 0..19 {
ds.union(i, i + 1);
}
let r = ds.find(0);
for i in 0..20 {
let fi = ds.find(i);
assert_eq!(fi, r);
assert_eq!(ds.find(fi), r);
}
assert_eq!(ds.count(), 1);
assert_eq!(ds.set_size(13), 20);
}
/// Depth without reading through `find`, which would compress the path
/// being measured.
fn depth_of(ds: &DisjointSet, mut x: usize) -> usize {
let mut d = 0;
while ds.parent[x] != x {
x = ds.parent[x];
d += 1;
}
d
}
/// Union by size bounds the tree height by log2(n) on its own, before any
/// path compression. The worst case for the bound is merging equal-sized
/// trees pairwise, which is what this builds: a balanced binary merge over
/// 2^14 elements, whose height must be at most 14.
#[test]
fn union_by_size_bounds_depth_by_log2() {
const K: usize = 14;
const N: usize = 1 << K;
let mut ds = DisjointSet::new(N);
let mut step = 1usize;
while step < N {
let mut i = 0usize;
while i + step < N {
// Both roots have exactly `step` elements here, so the tie
// rule decides and the height can grow by one per round.
assert!(ds.union(i, i + step));
i += 2 * step;
}
step *= 2;
}
assert_eq!(ds.count(), 1);
// Measure before any find(), which would compress the path being
// measured -- reading set_size(N - 1) here costs exactly the one
// deepest path and drops the observed height to K - 1.
let max_depth = (0..N).map(|i| depth_of(&ds, i)).max().unwrap();
assert!(
max_depth <= K,
"height {max_depth} exceeds the log2 bound {K}"
);
// It really does reach the bound, so this is not a vacuous assertion.
assert_eq!(max_depth, K);
assert_eq!(ds.set_size(N - 1), N);
// Path compression then flattens it: after one find() per node every
// node points straight at the root.
let r = ds.find(0);
for i in 0..N {
assert_eq!(ds.find(i), r);
}
for i in 0..N {
assert_eq!(depth_of(&ds, i), usize::from(i != r));
}
}
#[test]
fn empty_and_singleton() {
let mut e = DisjointSet::new(0);
assert!(e.is_empty());
assert_eq!(e.count(), 0);
assert!(e.sets().is_empty());
let mut s = DisjointSet::new(1);
assert!(!s.is_empty());
assert_eq!(s.count(), 1);
assert!(s.connected(0, 0));
assert!(!s.union(0, 0));
assert_eq!(s.count(), 1);
}
}