competitive_programming_rs/data_structure/
union_find.rs

1pub struct UnionFind {
2    parent: Vec<usize>,
3    sizes: Vec<usize>,
4    size: usize,
5}
6
7impl UnionFind {
8    pub fn new(n: usize) -> UnionFind {
9        UnionFind {
10            parent: (0..n).map(|i| i).collect::<Vec<usize>>(),
11            sizes: vec![1; n],
12            size: n,
13        }
14    }
15
16    pub fn find(&mut self, x: usize) -> usize {
17        if x == self.parent[x] {
18            x
19        } else {
20            let px = self.parent[x];
21            self.parent[x] = self.find(px);
22            self.parent[x]
23        }
24    }
25
26    pub fn unite(&mut self, x: usize, y: usize) -> bool {
27        let parent_x = self.find(x);
28        let parent_y = self.find(y);
29        if parent_x == parent_y {
30            return false;
31        }
32
33        let (large, small) = if self.sizes[parent_x] < self.sizes[parent_y] {
34            (parent_y, parent_x)
35        } else {
36            (parent_x, parent_y)
37        };
38
39        self.parent[small] = large;
40        self.sizes[large] += self.sizes[small];
41        self.sizes[small] = 0;
42        self.size -= 1;
43        true
44    }
45}
46
47#[cfg(test)]
48mod test {
49    use super::*;
50    use crate::utils::test_helper::Tester;
51
52    /// Solve http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=DSL_1_A
53    #[test]
54    fn solve_dsl_1_a() {
55        let tester = Tester::new("./assets/DSL_1_A/in/", "./assets/DSL_1_A/out/");
56
57        tester.test_solution(|sc| {
58            let n = sc.read();
59            let q = sc.read();
60            let mut uf = UnionFind::new(n);
61            for _ in 0..q {
62                let com: usize = sc.read();
63                let x = sc.read();
64                let y = sc.read();
65                if com == 0 {
66                    uf.unite(x, y);
67                } else {
68                    let ans = if uf.find(x) == uf.find(y) { 1 } else { 0 };
69                    sc.write(format!("{}\n", ans));
70                }
71            }
72        });
73    }
74}
competitive_programming_rs/data_structure/union_find.rs

competitive_programming_rs/data_structure/
union_find.rs
