Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Disjoint-Set / Union-Find - A simple
/*
UnionFind.rs
A simple single-file UnionFind/Disjoint-Set implementation.
AUTHOR: Joseph Catrambone <jo.jcat@gmail.com> (c) 2019
LICENSE: MIT
/// Example Usage:
///
/// ```
/// let uf = UnionFind::<char>::new();
/// uf.add('a');
/// uf.add('b');
/// uf.add('c');
/// assert_ne!(uf.find('a'), uf.find('b'));
/// assert_ne!(uf.find('b'), uf.find('c'));
/// uf.union('a', 'b');
/// assert_eq!(uf.find('a'), uf.find('b'));
/// assert_ne!(uf.find('b'), uf.find('c'));
/// ```
*/
extern crate hashbrown;
use hashbrown::HashMap;
use std::hash::Hash;
pub struct UnionFind<T> {
entry_parent_map: HashMap<T, usize>, // Position in Vec:: would suffice, but we don't want to search.
parent_parent_map: Vec<usize>, // len(p_p_map) gives us the next ID.
parent_node_count: Vec<usize>,
}
impl<T> UnionFind<T> where T: Eq + Hash {
pub fn new() -> Self {
UnionFind {
entry_parent_map: HashMap::<T, usize>::new(),
parent_parent_map: Vec::<usize>::new(),
parent_node_count: Vec::<usize>::new(),
}
}
/// Begin tracking a new node.
pub fn add(&mut self, node:T) {
assert!(!self.entry_parent_map.contains_key(&node));
let next_parent_id = self.parent_node_count.len();
self.entry_parent_map.insert(node, next_parent_id);
self.parent_parent_map.push(next_parent_id); // Own parent.
self.parent_node_count.push(1); // One item.
}
/// Join two disjoint nodes together into a common parent set.
pub fn union(&mut self, a:T, b:T) {
// Assign a and b to the same set.
let mut a_parent = self.find(a);
let mut b_parent = self.find(b);
if self.parent_node_count[a_parent] < self.parent_node_count[b_parent] {
self.parent_node_count[a_parent] += self.parent_node_count[b_parent];
self.parent_parent_map[b_parent] = a_parent;
} else {
self.parent_node_count[b_parent] += self.parent_node_count[a_parent];
self.parent_parent_map[a_parent] = b_parent;
}
}
/// Find the ID of the parent set of a node.
pub fn find(&mut self, a:T) -> usize {
let mut parent = *self.entry_parent_map.get(&a).expect("find called for element not yet added to tree.");
while self.parent_parent_map[parent] != parent {
parent = self.parent_parent_map[parent];
}
// Go back through the list and make all the assignments flat.
let new_parent_assignment = parent;
parent = *self.entry_parent_map.get(&a).unwrap();
while self.parent_parent_map[parent] != parent {
let next = self.parent_parent_map[parent];
self.parent_parent_map[parent] = new_parent_assignment;
parent = next;
}
new_parent_assignment
}
}
#[cfg(test)]
mod test {
use super::UnionFind;
#[test]
fn different_parents() {
let mut uf = UnionFind::<u32>::new();
uf.add(0);
uf.add(1);
assert_ne!(uf.find(0), uf.find(1));
}
#[test]
fn check_union() {
let mut uf = UnionFind::<u64>::new();
uf.add(420);
uf.add(69);
uf.union(420, 69); // Nice.
assert_eq!(uf.find(420), uf.find(69));
}
#[test]
fn check_speed() {
let mut uf = UnionFind::<u64>::new();
// It should be very fast to insert a million entries and check them as long as we're linear.
uf.add(0);
// With the naive approach, this takes hours. With the optimization, seconds.
for i in 1..1_000_000 {
uf.add(i);
uf.union(i-1, i); // Nice.
assert_eq!(uf.find(0), uf.find(i));
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.