// Copyright 2024 The Jujutsu Authors//// Licensed under the Apache License, Version 2.0 (the "License");// you may not use this file except in compliance with the License.// You may obtain a copy of the License at//// https://www.apache.org/licenses/LICENSE-2.0//// Unless required by applicable law or agreed to in writing, software// distributed under the License is distributed on an "AS IS" BASIS,// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.// See the License for the specific language governing permissions and// limitations under the License. //! This module implements a [`UnionFind<T>`] type which can be used to//! efficiently calculate disjoint sets for any data type. use std::collections::HashMap;use std::hash::Hash; #[derive(Clone, Copy)]struct Node<T> { root: T, size: u32,} /// Implementation of the union-find algorithm:/// <https://en.wikipedia.org/wiki/Disjoint-set_data_structure>////// Joins disjoint sets by size to amortize cost.#[derive(Clone)]pub struct UnionFind<T> { roots: HashMap<T, Node<T>>,} impl<T> Default for UnionFind<T>where T: Copy + Eq + Hash,{ fn default() -> Self { Self::new() }} impl<T> UnionFind<T>where T: Copy + Eq + Hash,{ /// Creates a new empty UnionFind data structure. pub fn new() -> Self { Self { roots: HashMap::new(), } } /// Returns the root identifying the union this item is a part of. pub fn find(&mut self, item: T) -> T { self.find_node(item).root } fn find_node(&mut self, item: T) -> Node<T> { match self.roots.get(&item) { Some(node) => { if node.root != item { let new_root = self.find_node(node.root); self.roots.insert(item, new_root); new_root } else { *node } } None => { let node = Node::<T> { root: item, size: 1, }; self.roots.insert(item, node); node } } } /// Unions the disjoint sets connected to `a` and `b`. pub fn union(&mut self, a: T, b: T) { let a = self.find_node(a); let b = self.find_node(b); if a.root == b.root { return; } let new_node = Node::<T> { root: if a.size < b.size { b.root } else { a.root }, size: a.size + b.size, }; self.roots.insert(a.root, new_node); self.roots.insert(b.root, new_node); }} #[cfg(test)]mod tests { use itertools::Itertools as _; use super::*; #[test] fn test_basic() { let mut union_find = UnionFind::<i32>::new(); // Everything starts as a singleton. assert_eq!(union_find.find(1), 1); assert_eq!(union_find.find(2), 2); assert_eq!(union_find.find(3), 3); // Make two pair sets. This implicitly adds node 4. union_find.union(1, 2); union_find.union(3, 4); assert_eq!(union_find.find(1), union_find.find(2)); assert_eq!(union_find.find(3), union_find.find(4)); assert_ne!(union_find.find(1), union_find.find(3)); // Unioning the pairs gives everything the same root. union_find.union(1, 3); assert!( [ union_find.find(1), union_find.find(2), union_find.find(3), union_find.find(4), ] .iter() .all_equal() ); } #[test] fn test_union_by_size() { let mut union_find = UnionFind::<i32>::new(); // Create a set of 3 and a set of 2. union_find.union(1, 2); union_find.union(2, 3); union_find.union(4, 5); let set3 = union_find.find(1); let set2 = union_find.find(4); assert_ne!(set3, set2); // Merging them always chooses the larger set. let mut large_first = union_find.clone(); large_first.union(1, 4); assert_eq!(large_first.find(1), set3); assert_eq!(large_first.find(4), set3); let mut small_first = union_find.clone(); small_first.union(4, 1); assert_eq!(small_first.find(1), set3); assert_eq!(small_first.find(4), set3); }}
// Copyright 2024 The Jujutsu Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//! This module implements a [`UnionFind<T>`] type which can be used to
//! efficiently calculate disjoint sets for any data type.
use std::collections::HashMap;
use std::hash::Hash;
#[derive(Clone, Copy)]
struct Node<T> {
root: T,
size: u32,
}
/// Implementation of the union-find algorithm:
/// <https://en.wikipedia.org/wiki/Disjoint-set_data_structure>
///
/// Joins disjoint sets by size to amortize cost.
#[derive(Clone)]
pub struct UnionFind<T> {
roots: HashMap<T, Node<T>>,
}
impl<T> Default for UnionFind<T>
where
T: Copy + Eq + Hash,
{
fn default() -> Self {
Self::new()
}
}
impl<T> UnionFind<T>
where
T: Copy + Eq + Hash,
{
/// Creates a new empty UnionFind data structure.
pub fn new() -> Self {
Self {
roots: HashMap::new(),
}
}
/// Returns the root identifying the union this item is a part of.
pub fn find(&mut self, item: T) -> T {
self.find_node(item).root
}
fn find_node(&mut self, item: T) -> Node<T> {
match self.roots.get(&item) {
Some(node) => {
if node.root != item {
let new_root = self.find_node(node.root);
self.roots.insert(item, new_root);
new_root
} else {
*node
}
}
None => {
let node = Node::<T> {
root: item,
size: 1,
};
self.roots.insert(item, node);
node
}
}
}
/// Unions the disjoint sets connected to `a` and `b`.
pub fn union(&mut self, a: T, b: T) {
let a = self.find_node(a);
let b = self.find_node(b);
if a.root == b.root {
return;
}
let new_node = Node::<T> {
root: if a.size < b.size { b.root } else { a.root },
size: a.size + b.size,
};
self.roots.insert(a.root, new_node);
self.roots.insert(b.root, new_node);
}
}
#[cfg(test)]
mod tests {
use itertools::Itertools as _;
use super::*;
#[test]
fn test_basic() {
let mut union_find = UnionFind::<i32>::new();
// Everything starts as a singleton.
assert_eq!(union_find.find(1), 1);
assert_eq!(union_find.find(2), 2);
assert_eq!(union_find.find(3), 3);
// Make two pair sets. This implicitly adds node 4.
union_find.union(1, 2);
union_find.union(3, 4);
assert_eq!(union_find.find(1), union_find.find(2));
assert_eq!(union_find.find(3), union_find.find(4));
assert_ne!(union_find.find(1), union_find.find(3));
// Unioning the pairs gives everything the same root.
union_find.union(1, 3);
assert!(
[
union_find.find(1),
union_find.find(2),
union_find.find(3),
union_find.find(4),
]
.iter()
.all_equal()
);
}
#[test]
fn test_union_by_size() {
let mut union_find = UnionFind::<i32>::new();
// Create a set of 3 and a set of 2.
union_find.union(1, 2);
union_find.union(2, 3);
union_find.union(4, 5);
let set3 = union_find.find(1);
let set2 = union_find.find(4);
assert_ne!(set3, set2);
// Merging them always chooses the larger set.
let mut large_first = union_find.clone();
large_first.union(1, 4);
assert_eq!(large_first.find(1), set3);
assert_eq!(large_first.find(4), set3);
let mut small_first = union_find.clone();
small_first.union(4, 1);
assert_eq!(small_first.find(1), set3);
assert_eq!(small_first.find(4), set3);
}
}