#![allow(non_snake_case)]
#![allow(non_camel_case_types)]
use std::marker::PhantomData;
pub trait ToFromU32<T: Sized = Self> {
fn to_u32(x: Self) -> u32;
fn from_u32(x: u32) -> Self;
}
impl ToFromU32 for u32 {
fn to_u32(x: u32) -> u32 {
x
}
fn from_u32(x: u32) -> u32 {
x
}
}
pub struct UnionFind<T: ToFromU32> {
parent_or_size: Vec<i32>,
anchor: PhantomData<T>,
}
const UF_MAX_SIZE: u32 = 0x7FFF_FFF0;
impl<T: ToFromU32> UnionFind<T> {
pub fn new(size: usize) -> Self {
if size > UF_MAX_SIZE as usize {
panic!("UnionFind::new: too many elements; max = 2^31 - 16.");
}
let mut parent_or_size = Vec::<i32>::new();
parent_or_size.resize(size, -1);
Self {
parent_or_size,
anchor: PhantomData,
}
}
#[inline(always)]
fn find(&mut self, elem0: u32) -> u32 {
let elem0_parent_or_size: i32 = self.parent_or_size[elem0 as usize];
if elem0_parent_or_size < 0 {
return elem0;
}
let elem1 = elem0_parent_or_size as u32;
let elem1_parent_or_size: i32 = self.parent_or_size[elem1 as usize];
if elem1_parent_or_size < 0 {
self.parent_or_size[elem0 as usize] = elem1 as i32;
return elem1;
}
let elem2 = elem1_parent_or_size as u32;
let elem2_parent_or_size: i32 = self.parent_or_size[elem2 as usize];
if elem2_parent_or_size < 0 {
self.parent_or_size[elem1 as usize] = elem2 as i32;
self.parent_or_size[elem0 as usize] = elem2 as i32;
return elem2;
}
let elem3 = elem2_parent_or_size as u32;
let elem3_parent_or_size: i32 = self.parent_or_size[elem3 as usize];
if elem3_parent_or_size < 0 {
self.parent_or_size[elem2 as usize] = elem3 as i32;
self.parent_or_size[elem1 as usize] = elem3 as i32;
self.parent_or_size[elem0 as usize] = elem3 as i32;
return elem3;
}
let elem4 = elem3_parent_or_size as u32;
let root = self.find_slow(elem4);
assert!(root < UF_MAX_SIZE);
self.parent_or_size[elem3 as usize] = root as i32;
self.parent_or_size[elem2 as usize] = root as i32;
self.parent_or_size[elem1 as usize] = root as i32;
self.parent_or_size[elem0 as usize] = root as i32;
return root;
}
#[inline(never)]
fn find_slow(&mut self, elem0: u32) -> u32 {
let elem0_parent_or_size: i32 = self.parent_or_size[elem0 as usize];
if elem0_parent_or_size < 0 {
return elem0;
}
let elem1 = elem0_parent_or_size as u32;
let elem1_parent_or_size: i32 = self.parent_or_size[elem1 as usize];
if elem1_parent_or_size < 0 {
self.parent_or_size[elem0 as usize] = elem1 as i32;
return elem1;
}
let elem2 = elem1_parent_or_size as u32;
let root = self.find_slow(elem2);
assert!(root < UF_MAX_SIZE);
self.parent_or_size[elem1 as usize] = root as i32;
self.parent_or_size[elem0 as usize] = root as i32;
return root;
}
pub fn union(&mut self, elem1t: T, elem2t: T) {
let elem1 = ToFromU32::to_u32(elem1t);
let elem2 = ToFromU32::to_u32(elem2t);
if elem1 == elem2 {
return;
}
let root1: u32 = self.find(elem1);
let root2: u32 = self.find(elem2);
if root1 == root2 {
return;
}
let size1: i32 = self.parent_or_size[root1 as usize];
let size2: i32 = self.parent_or_size[root2 as usize];
assert!(size1 < 0 && size2 < 0);
if size1 < size2 {
self.parent_or_size[root1 as usize] = root2 as i32;
self.parent_or_size[root2 as usize] += size1;
} else {
self.parent_or_size[root2 as usize] = root1 as i32;
self.parent_or_size[root1 as usize] += size2;
}
}
}
const UFEC_NULL: u32 = 0xFFFF_FFFF;
#[derive(Clone)]
struct LLElem {
elem: u32,
tail: u32,
}
pub struct UnionFindEquivClasses<T: ToFromU32> {
heads: Vec<u32>,
lists: Vec<LLElem>,
anchor: PhantomData<T>,
}
impl<T: ToFromU32> UnionFind<T> {
pub fn get_equiv_classes(&mut self) -> UnionFindEquivClasses<T> {
let nElemsUSize = self.parent_or_size.len();
assert!(nElemsUSize < UF_MAX_SIZE as usize);
let nElems = nElemsUSize as u32;
let mut heads = Vec::<u32>::new();
heads.resize(nElems as usize, UFEC_NULL);
let mut lists = Vec::<LLElem>::new();
lists.resize(
nElems as usize,
LLElem {
elem: 0,
tail: UFEC_NULL,
},
);
for i in 0..nElems {
if self.parent_or_size[i as usize] >= 0 {
let root_i: u32 = self.find(i);
assert!(root_i < 0x8000_0000u32);
heads[i as usize] = root_i;
}
}
let mut list_bump = 0u32;
for i in 0..nElems {
if self.parent_or_size[i as usize] < 0 {
lists[list_bump as usize] = LLElem {
elem: i,
tail: if heads[i as usize] == UFEC_NULL {
UFEC_NULL
} else {
heads[i as usize] & 0x7FFF_FFFF
},
};
assert!(list_bump < 0x8000_0000u32);
heads[i as usize] = list_bump | 0x8000_0000u32;
list_bump += 1;
} else {
let i_root = heads[i as usize];
lists[list_bump as usize] = LLElem {
elem: i,
tail: if heads[i_root as usize] == UFEC_NULL {
UFEC_NULL
} else {
heads[i_root as usize] & 0x7FFF_FFFF
},
};
assert!(list_bump < 0x8000_0000u32);
heads[i_root as usize] = list_bump | 0x8000_0000u32;
list_bump += 1;
}
}
assert!(list_bump == nElems);
assert!(heads.len() == nElemsUSize);
assert!(lists.len() == nElemsUSize);
UnionFindEquivClasses {
heads,
lists,
anchor: PhantomData,
}
}
}
impl<T: ToFromU32> UnionFindEquivClasses<T> {
pub fn in_same_equivalence_class(&self, item1: T, item2: T) -> Option<bool> {
let mut item1num = ToFromU32::to_u32(item1) as usize;
let mut item2num = ToFromU32::to_u32(item2) as usize;
if item1num >= self.heads.len() || item2num >= self.heads.len() {
return None;
}
if (self.heads[item1num] & 0x8000_0000) == 0 {
item1num = self.heads[item1num] as usize;
}
if (self.heads[item2num] & 0x8000_0000) == 0 {
item2num = self.heads[item2num] as usize;
}
debug_assert!((self.heads[item1num] & 0x8000_0000) == 0x8000_0000);
debug_assert!((self.heads[item2num] & 0x8000_0000) == 0x8000_0000);
Some(item1num == item2num)
}
}
pub struct UnionFindEquivClassElemsIter<'a, T: ToFromU32> {
ufec: &'a UnionFindEquivClasses<T>,
next: u32,
}
impl<T: ToFromU32> UnionFindEquivClasses<T> {
pub fn equiv_class_elems_iter<'a>(&'a self, item: T) -> UnionFindEquivClassElemsIter<'a, T> {
let mut itemU32 = ToFromU32::to_u32(item);
assert!((itemU32 as usize) < self.heads.len());
if (self.heads[itemU32 as usize] & 0x8000_0000) == 0 {
itemU32 = self.heads[itemU32 as usize];
}
assert!((self.heads[itemU32 as usize] & 0x8000_0000) == 0x8000_0000);
let next = self.heads[itemU32 as usize] & 0x7FFF_FFFF;
UnionFindEquivClassElemsIter { ufec: &self, next }
}
}
impl<'a, T: ToFromU32> Iterator for UnionFindEquivClassElemsIter<'a, T> {
type Item = T;
fn next(&mut self) -> Option<Self::Item> {
if self.next == UFEC_NULL {
None
} else {
let res: T = ToFromU32::from_u32(self.ufec.lists[self.next as usize].elem);
self.next = self.ufec.lists[self.next as usize].tail;
Some(res)
}
}
}
pub struct UnionFindEquivClassLeadersIter<'a, T: ToFromU32> {
ufec: &'a UnionFindEquivClasses<T>,
next: u32,
}
impl<T: ToFromU32> UnionFindEquivClasses<T> {
pub fn equiv_class_leaders_iter<'a>(&'a self) -> UnionFindEquivClassLeadersIter<'a, T> {
UnionFindEquivClassLeadersIter {
ufec: &self,
next: 0,
}
}
}
impl<'a, T: ToFromU32> Iterator for UnionFindEquivClassLeadersIter<'a, T> {
type Item = T;
fn next(&mut self) -> Option<Self::Item> {
loop {
if self.next as usize >= self.ufec.heads.len() {
return None;
}
if (self.ufec.heads[self.next as usize] & 0x8000_0000) == 0x8000_0000 {
let res = ToFromU32::from_u32(self.next);
self.next += 1;
return Some(res);
}
self.next += 1;
}
}
}
#[cfg(test)]
mod union_find_test_utils {
use super::UnionFindEquivClasses;
pub fn test_eclass(eclasses: &UnionFindEquivClasses<u32>, elem: u32, expected: &Vec<u32>) {
let mut expected_sorted = expected.clone();
let mut actual = vec![];
for ecm in eclasses.equiv_class_elems_iter(elem) {
actual.push(ecm);
}
expected_sorted.sort();
actual.sort();
assert!(actual == expected_sorted);
}
pub fn test_leaders(
univ_size: u32,
eclasses: &UnionFindEquivClasses<u32>,
expected: &Vec<u32>,
) {
let mut actual = vec![];
for leader in eclasses.equiv_class_leaders_iter() {
actual.push(leader);
}
assert!(actual == *expected);
let mut univ_actual = vec![];
for leader in eclasses.equiv_class_leaders_iter() {
for elem in eclasses.equiv_class_elems_iter(leader) {
univ_actual.push(elem);
}
}
univ_actual.sort();
let mut univ_expected = vec![];
for i in 0..univ_size {
univ_expected.push(i);
}
assert!(univ_actual == univ_expected);
}
pub fn test_in_same_eclass(
eclasses: &UnionFindEquivClasses<u32>,
elem1: u32,
elem2: u32,
expected: Option<bool>,
) {
assert!(eclasses.in_same_equivalence_class(elem1, elem2) == expected);
assert!(eclasses.in_same_equivalence_class(elem2, elem1) == expected);
}
}
#[test]
fn test_union_find() {
const UNIV_SIZE: u32 = 8;
let mut uf = UnionFind::new(UNIV_SIZE as usize);
let mut uf_eclasses = uf.get_equiv_classes();
union_find_test_utils::test_eclass(&uf_eclasses, 0, &vec![0]);
union_find_test_utils::test_eclass(&uf_eclasses, 1, &vec![1]);
union_find_test_utils::test_eclass(&uf_eclasses, 2, &vec![2]);
union_find_test_utils::test_eclass(&uf_eclasses, 3, &vec![3]);
union_find_test_utils::test_eclass(&uf_eclasses, 4, &vec![4]);
union_find_test_utils::test_eclass(&uf_eclasses, 5, &vec![5]);
union_find_test_utils::test_eclass(&uf_eclasses, 6, &vec![6]);
union_find_test_utils::test_eclass(&uf_eclasses, 7, &vec![7]);
union_find_test_utils::test_leaders(UNIV_SIZE, &uf_eclasses, &vec![0, 1, 2, 3, 4, 5, 6, 7]);
uf.union(2, 4);
uf_eclasses = uf.get_equiv_classes();
union_find_test_utils::test_eclass(&uf_eclasses, 0, &vec![0]);
union_find_test_utils::test_eclass(&uf_eclasses, 1, &vec![1]);
union_find_test_utils::test_eclass(&uf_eclasses, 2, &vec![4, 2]);
union_find_test_utils::test_eclass(&uf_eclasses, 3, &vec![3]);
union_find_test_utils::test_eclass(&uf_eclasses, 4, &vec![4, 2]);
union_find_test_utils::test_eclass(&uf_eclasses, 5, &vec![5]);
union_find_test_utils::test_eclass(&uf_eclasses, 6, &vec![6]);
union_find_test_utils::test_eclass(&uf_eclasses, 7, &vec![7]);
union_find_test_utils::test_leaders(UNIV_SIZE, &uf_eclasses, &vec![0, 1, 2, 3, 5, 6, 7]);
uf.union(5, 3);
uf_eclasses = uf.get_equiv_classes();
union_find_test_utils::test_eclass(&uf_eclasses, 0, &vec![0]);
union_find_test_utils::test_eclass(&uf_eclasses, 1, &vec![1]);
union_find_test_utils::test_eclass(&uf_eclasses, 2, &vec![4, 2]);
union_find_test_utils::test_eclass(&uf_eclasses, 3, &vec![5, 3]);
union_find_test_utils::test_eclass(&uf_eclasses, 4, &vec![4, 2]);
union_find_test_utils::test_eclass(&uf_eclasses, 5, &vec![5, 3]);
union_find_test_utils::test_eclass(&uf_eclasses, 6, &vec![6]);
union_find_test_utils::test_eclass(&uf_eclasses, 7, &vec![7]);
union_find_test_utils::test_leaders(UNIV_SIZE, &uf_eclasses, &vec![0, 1, 2, 5, 6, 7]);
uf.union(2, 5);
uf_eclasses = uf.get_equiv_classes();
union_find_test_utils::test_eclass(&uf_eclasses, 0, &vec![0]);
union_find_test_utils::test_eclass(&uf_eclasses, 1, &vec![1]);
union_find_test_utils::test_eclass(&uf_eclasses, 2, &vec![5, 4, 3, 2]);
union_find_test_utils::test_eclass(&uf_eclasses, 3, &vec![5, 4, 3, 2]);
union_find_test_utils::test_eclass(&uf_eclasses, 4, &vec![5, 4, 3, 2]);
union_find_test_utils::test_eclass(&uf_eclasses, 5, &vec![5, 4, 3, 2]);
union_find_test_utils::test_eclass(&uf_eclasses, 6, &vec![6]);
union_find_test_utils::test_eclass(&uf_eclasses, 7, &vec![7]);
union_find_test_utils::test_leaders(UNIV_SIZE, &uf_eclasses, &vec![0, 1, 2, 6, 7]);
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 0, 0, Some(true));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 0, 1, Some(false));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 0, 2, Some(false));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 0, 3, Some(false));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 0, 4, Some(false));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 0, 5, Some(false));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 0, 6, Some(false));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 0, 7, Some(false));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 1, 1, Some(true));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 1, 2, Some(false));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 1, 3, Some(false));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 1, 4, Some(false));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 1, 5, Some(false));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 1, 6, Some(false));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 1, 7, Some(false));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 2, 2, Some(true));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 2, 3, Some(true));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 2, 4, Some(true));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 2, 5, Some(true));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 2, 6, Some(false));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 2, 7, Some(false));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 3, 3, Some(true));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 3, 4, Some(true));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 3, 5, Some(true));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 3, 6, Some(false));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 3, 7, Some(false));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 4, 4, Some(true));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 4, 5, Some(true));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 4, 6, Some(false));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 4, 7, Some(false));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 5, 5, Some(true));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 5, 6, Some(false));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 5, 7, Some(false));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 6, 6, Some(true));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 6, 7, Some(false));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 7, 7, Some(true));
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 0, 8, None);
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 8, 0, None);
union_find_test_utils::test_in_same_eclass(&uf_eclasses, 8, 8, None);
uf.union(7, 1);
uf_eclasses = uf.get_equiv_classes();
union_find_test_utils::test_eclass(&uf_eclasses, 0, &vec![0]);
union_find_test_utils::test_eclass(&uf_eclasses, 1, &vec![7, 1]);
union_find_test_utils::test_eclass(&uf_eclasses, 2, &vec![5, 4, 3, 2]);
union_find_test_utils::test_eclass(&uf_eclasses, 3, &vec![5, 4, 3, 2]);
union_find_test_utils::test_eclass(&uf_eclasses, 4, &vec![5, 4, 3, 2]);
union_find_test_utils::test_eclass(&uf_eclasses, 5, &vec![5, 4, 3, 2]);
union_find_test_utils::test_eclass(&uf_eclasses, 6, &vec![6]);
union_find_test_utils::test_eclass(&uf_eclasses, 7, &vec![7, 1]);
union_find_test_utils::test_leaders(UNIV_SIZE, &uf_eclasses, &vec![0, 2, 6, 7]);
uf.union(6, 7);
uf_eclasses = uf.get_equiv_classes();
union_find_test_utils::test_eclass(&uf_eclasses, 0, &vec![0]);
union_find_test_utils::test_eclass(&uf_eclasses, 1, &vec![7, 6, 1]);
union_find_test_utils::test_eclass(&uf_eclasses, 2, &vec![5, 4, 3, 2]);
union_find_test_utils::test_eclass(&uf_eclasses, 3, &vec![5, 4, 3, 2]);
union_find_test_utils::test_eclass(&uf_eclasses, 4, &vec![5, 4, 3, 2]);
union_find_test_utils::test_eclass(&uf_eclasses, 5, &vec![5, 4, 3, 2]);
union_find_test_utils::test_eclass(&uf_eclasses, 6, &vec![7, 6, 1]);
union_find_test_utils::test_eclass(&uf_eclasses, 7, &vec![7, 6, 1]);
union_find_test_utils::test_leaders(UNIV_SIZE, &uf_eclasses, &vec![0, 2, 6]);
uf.union(4, 1);
uf_eclasses = uf.get_equiv_classes();
union_find_test_utils::test_eclass(&uf_eclasses, 0, &vec![0]);
union_find_test_utils::test_eclass(&uf_eclasses, 1, &vec![7, 6, 5, 4, 3, 2, 1]);
union_find_test_utils::test_eclass(&uf_eclasses, 2, &vec![7, 6, 5, 4, 3, 2, 1]);
union_find_test_utils::test_eclass(&uf_eclasses, 3, &vec![7, 6, 5, 4, 3, 2, 1]);
union_find_test_utils::test_eclass(&uf_eclasses, 4, &vec![7, 6, 5, 4, 3, 2, 1]);
union_find_test_utils::test_eclass(&uf_eclasses, 5, &vec![7, 6, 5, 4, 3, 2, 1]);
union_find_test_utils::test_eclass(&uf_eclasses, 6, &vec![7, 6, 5, 4, 3, 2, 1]);
union_find_test_utils::test_eclass(&uf_eclasses, 7, &vec![7, 6, 5, 4, 3, 2, 1]);
union_find_test_utils::test_leaders(UNIV_SIZE, &uf_eclasses, &vec![0, 6]);
uf.union(0, 3);
uf_eclasses = uf.get_equiv_classes();
union_find_test_utils::test_eclass(&uf_eclasses, 0, &vec![7, 6, 5, 4, 3, 2, 1, 0]);
union_find_test_utils::test_eclass(&uf_eclasses, 1, &vec![7, 6, 5, 4, 3, 2, 1, 0]);
union_find_test_utils::test_eclass(&uf_eclasses, 2, &vec![7, 6, 5, 4, 3, 2, 1, 0]);
union_find_test_utils::test_eclass(&uf_eclasses, 3, &vec![7, 6, 5, 4, 3, 2, 1, 0]);
union_find_test_utils::test_eclass(&uf_eclasses, 4, &vec![7, 6, 5, 4, 3, 2, 1, 0]);
union_find_test_utils::test_eclass(&uf_eclasses, 5, &vec![7, 6, 5, 4, 3, 2, 1, 0]);
union_find_test_utils::test_eclass(&uf_eclasses, 6, &vec![7, 6, 5, 4, 3, 2, 1, 0]);
union_find_test_utils::test_eclass(&uf_eclasses, 7, &vec![7, 6, 5, 4, 3, 2, 1, 0]);
union_find_test_utils::test_leaders(UNIV_SIZE, &uf_eclasses, &vec![0]);
uf.union(1, 2);
uf_eclasses = uf.get_equiv_classes();
union_find_test_utils::test_eclass(&uf_eclasses, 0, &vec![7, 6, 5, 4, 3, 2, 1, 0]);
union_find_test_utils::test_eclass(&uf_eclasses, 1, &vec![7, 6, 5, 4, 3, 2, 1, 0]);
union_find_test_utils::test_eclass(&uf_eclasses, 2, &vec![7, 6, 5, 4, 3, 2, 1, 0]);
union_find_test_utils::test_eclass(&uf_eclasses, 3, &vec![7, 6, 5, 4, 3, 2, 1, 0]);
union_find_test_utils::test_eclass(&uf_eclasses, 4, &vec![7, 6, 5, 4, 3, 2, 1, 0]);
union_find_test_utils::test_eclass(&uf_eclasses, 5, &vec![7, 6, 5, 4, 3, 2, 1, 0]);
union_find_test_utils::test_eclass(&uf_eclasses, 6, &vec![7, 6, 5, 4, 3, 2, 1, 0]);
union_find_test_utils::test_eclass(&uf_eclasses, 7, &vec![7, 6, 5, 4, 3, 2, 1, 0]);
union_find_test_utils::test_leaders(UNIV_SIZE, &uf_eclasses, &vec![0]);
}