atcoder_union_find.rs

mod union_find{
    pub struct UnionFind{
        pub parent: Vec<usize>,
        rank: Vec<usize>,
        islands: i64
    }

    impl UnionFind{
        pub fn new(n: usize) -> UnionFind{
            return UnionFind{parent: (0..n).collect(),
                rank: vec![0; n],
                islands: n as i64};
        }

        pub fn find(&mut self, a: usize) -> usize{
            if self.parent[a] == a{
                return a;
            }else{
                self.parent[a] = self.find(self.parent[a]);
                return self.parent[a];
            }
        }

        // あらたにつながったならtrue
        pub fn unite(&mut self, a: usize, b: usize) -> bool{
            assert!(a < self.parent.len());
            assert!(b < self.parent.len());
            let a = self.find(a);
            let b = self.find(b);
            if a == b{
                return false;
            }

            if self.rank[a] < self.rank[b]{
                self.parent[a] = b;
            }else{
                self.parent[b] = a;
                if self.rank[a] == self.rank[b]{
                    self.rank[a] += 1;
                }
            }
            return true;
        }

        pub fn same(&mut self, a: usize, b: usize) -> bool{
            return self.find(a) == self.find(b)
        }

    }
}