Rust Disjoint Set Union (DSU)

dsu.rs

mod dsu {
    pub struct dsu {
        parents: Vec<usize>,
        ranks: Vec<usize>,
    }

    impl dsu {
        pub fn new(size: usize) -> Self {
            Self {
                parents: (0..size).collect(),
                ranks: vec![1; size],
            }
        }

        pub fn is_valid(&mut self, id: usize) -> bool {
            id >= 0 && id < self.parents.len()
        }

        pub fn find(&mut self, id: usize) -> usize {
            if id != self.parents[id] {
                self.parents[id] = self.find(self.parents[id]);
            }
            self.parents[id]
        }

        pub fn same(&mut self, mut id1: usize, mut id2: usize) -> bool {
            id1 = self.find(id1);
            id2 = self.find(id2);
            id1 == id2
        }

        pub fn union(&mut self, mut id1: usize, mut id2: usize) -> bool {
            id1 = self.find(id1);
            id2 = self.find(id2);

            if id1 == id2 {
                return false;
            }

            if self.ranks[id1] < self.ranks[id2] {
                std::mem::swap(&mut id2, &mut id1);
            }

            self.parents[id2] = id1;
            self.ranks[id1] += self.ranks[id2];

            true
        }
    }
}