union_find/
lib.rs

1//! Generic Union-Find (Disjoint Set) data structure.
2//!
3//! Each component carries:
4//! - a `size` (used for union-by-size),
5//! - a `character` value (an arbitrary `usize` updated through the user-supplied `op` closure).
6//!
7//! `union(x, y, w)` merges the components of `x` and `y`, applying the closure
8//! `op` to combine their characters along with the edge weight `w`. When `x` and
9//! `y` are already in the same component, only `character[root] = op(character[root], w)`
10//! is updated and the function returns `false`; otherwise it returns `true`.
11
12/// Generic Union-Find with weighted union and per-component character.
13pub struct UnionFind<F>
14where
15    F: Fn(usize, usize) -> usize,
16{
17    root: Vec<usize>,
18    size: Vec<usize>,
19    components: usize,
20    character: Vec<usize>,
21    op: F,
22}
23
24impl<F> UnionFind<F>
25where
26    F: Fn(usize, usize) -> usize,
27{
28    /// Build a Union-Find over the indices `0..=n`. The closure `op` defines how
29    /// character values combine during `union`. `init` seeds every component's
30    /// character.
31    pub fn new(n: usize, op: F, init: usize) -> Self {
32        Self {
33            root: (0..=n).collect(),
34            size: vec![1; n + 1],
35            components: n,
36            character: vec![init; n + 1],
37            op,
38        }
39    }
40
41    /// Find the representative of `x` with path compression.
42    pub fn find(&mut self, x: usize) -> usize {
43        if self.root[x] == x {
44            return x;
45        }
46        self.root[x] = self.find(self.root[x]);
47        self.root[x]
48    }
49
50    /// Union the components of `x` and `y`, combining characters via `op` and `w`.
51    /// Returns `true` if a merge happened, `false` if `x` and `y` were already
52    /// in the same component.
53    pub fn union(&mut self, x: usize, y: usize, w: usize) -> bool {
54        let root_x = self.find(x);
55        let root_y = self.find(y);
56
57        if root_x == root_y {
58            self.character[root_x] = (self.op)(self.character[root_x], w);
59            return false;
60        }
61
62        if self.size[root_x] > self.size[root_y] {
63            self.size[root_x] += self.size[root_y];
64            self.character[root_x] =
65                (self.op)((self.op)(self.character[root_x], self.character[root_y]), w);
66            self.root[root_y] = root_x;
67        } else {
68            self.size[root_y] += self.size[root_x];
69            self.character[root_y] =
70                (self.op)((self.op)(self.character[root_x], self.character[root_y]), w);
71            self.root[root_x] = root_y;
72        }
73        self.components -= 1;
74        true
75    }
76
77    /// Whether `x` and `y` are in the same component.
78    pub fn are_connected(&mut self, x: usize, y: usize) -> bool {
79        self.find(x) == self.find(y)
80    }
81
82    /// Number of distinct components remaining.
83    pub fn get_components(&self) -> usize {
84        self.components
85    }
86
87    /// Representative of `x`.
88    pub fn get_root(&mut self, x: usize) -> usize {
89        self.find(x)
90    }
91
92    /// Character value associated with the component containing `x`.
93    pub fn get_character(&mut self, x: usize) -> usize {
94        let root_x = self.find(x);
95        self.character[root_x]
96    }
97
98    /// Size of the component containing `x`.
99    pub fn get_size(&mut self, x: usize) -> usize {
100        let root_x = self.find(x);
101        self.size[root_x]
102    }
103}
union_find/lib.rs

union_find/
lib.rs
