competitive_programming_lib/DataStructures/
union_find.rs

1pub struct UnionFind {
2    parent: Vec<usize>,
3    rank: Vec<usize>,
4}
5
6impl UnionFind {
7    pub fn new(size: usize) -> Self {
8        UnionFind {
9            parent: (0..size).collect(),
10            rank: vec![0; size],
11        }
12    }
13
14    pub fn find(&mut self, mut x: usize) -> usize {
15        if self.parent[x] != x {
16            self.parent[x] = self.find(self.parent[x]);
17        }
18        self.parent[x]
19    }
20
21    pub fn union(&mut self, x: usize, y: usize) {
22        let root_x = self.find(x);
23        let root_y = self.find(y);
24
25        if root_x != root_y {
26            if self.rank[root_x] > self.rank[root_y] {
27                self.parent[root_y] = root_x;
28            } else if self.rank[root_x] < self.rank[root_y] {
29                self.parent[root_x] = root_y;
30            } else {
31                self.parent[root_y] = root_x;
32                self.rank[root_x] += 1;
33            }
34        }
35    }
36}