Disjoint-Set / Union-Find - A simple

unionfind.rs

/*
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));
		}
	}
}