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
//! Disjoint set forest over dense indices, used to join topology nodes into
//! calculated buses.
pub(crate) struct UnionFind {
parent: Vec<usize>,
}
impl UnionFind {
pub(crate) fn new(len: usize) -> Self {
Self {
parent: (0..len).collect(),
}
}
/// The representative of `value`'s set, compressing the visited chain.
///
/// The walk is iterative. A source lists its switches in whatever order it
/// likes, and joining each node to the one before it builds a chain as
/// long as the node count, so recursion here would exhaust the stack on a
/// large node breaker model.
pub(crate) fn find(&mut self, value: usize) -> usize {
let mut root = value;
while self.parent[root] != root {
root = self.parent[root];
}
let mut walk = value;
while self.parent[walk] != root {
let next = self.parent[walk];
self.parent[walk] = root;
walk = next;
}
root
}
/// Merge the sets containing `first` and `second`; the representative of
/// `first`'s set stays the representative. Callers order their calculated
/// buses by that representative, so the rule is part of the decoded
/// result and not an implementation detail.
pub(crate) fn union(&mut self, first: usize, second: usize) {
let first = self.find(first);
let second = self.find(second);
if first != second {
self.parent[second] = first;
}
}
}
#[cfg(test)]
mod tests {
use super::UnionFind;
#[test]
fn a_long_chain_of_joins_resolves_without_recursion() {
const LEN: usize = 500_000;
let mut union = UnionFind::new(LEN);
// Joining each new root to the preceding node in increasing order
// grows the parent chain by one on every call.
for value in 0..LEN - 1 {
union.union(value + 1, value);
}
let root = union.find(0);
assert!((0..LEN).all(|value| union.find(value) == root));
}
#[test]
fn the_first_set_keeps_its_representative() {
let mut union = UnionFind::new(6);
union.union(3, 1);
assert_eq!(union.find(1), 3);
union.union(2, 3);
assert_eq!(union.find(1), 2);
assert_ne!(union.find(4), union.find(5));
}
}