使用 Rust 实现并查集

最近在刷 Leetcode, 每日题打卡题经常出现并查集, 这里记录下使用 Rust 实现的并查集.

    pub struct UnionFind {
        n: usize,
        fa: Vec<usize>,
    }

    impl UnionFind {
        pub fn new(n: usize) -> Self {
            Self {
                n,
                fa: (0..n).collect(),
            }
        }

        // 并查集中不同集合的数量
        pub fn count(&self) -> usize {
            let mut c = 0;
            for i in 0..self.fa.len() {
                if self.fa[i] == i {
                    c += 1;
                }
            }
            c
        }

        pub fn find(&mut self, x: usize) -> usize {
            assert_eq!(x < self.n, true);
            if self.fa[x] == x {
                x
            } else {
                // 路径压缩
                let find_fa = self.find(self.fa[x]);
                self.fa[x] = find_fa;
                self.fa[x]
            }
        }

        pub fn merge(&mut self, i: usize, j: usize) {
            let ifa = self.find(i);
            let jfa = self.find(j);
            self.fa[ifa] = jfa;
        }

        pub fn connected(&mut self, i: usize, j: usize) -> bool {
            self.find(i) == self.find(j)
        }
    }