graph_tools/
union_find.rs

1pub struct UnionFind{
2    n_components: usize,
3    pub parent: Vec<usize>
4}
5
6impl UnionFind{
7    pub fn new(n : usize) -> Self {
8        Self {
9            n_components: n,
10            parent: (0..n).collect(),
11        }
12    }
13
14    pub fn reset<'a>(&mut self, nodes : impl Iterator<Item = &'a usize>) {
15        let mut n_nodes = 0;
16        for n in nodes{
17            self.parent[*n] = *n;
18            n_nodes += 1;
19        }
20        self.n_components = n_nodes;
21    }
22
23    pub fn union(&mut self, mut u: usize, mut v: usize){
24        let p = &mut self.parent;
25        
26        while p[u] != p[v] {
27            if p[u] > p[v] {
28                if u == p[u] {
29                    p[u] = p[v];
30                    self.n_components-=1;
31                    return;
32                }
33                let z = p[u];
34                p[u] = p[v];
35                u = z;
36            }
37            else {
38                if v == p[v] {
39                    p[v] = p[u];
40                    self.n_components-=1;
41                    return;
42                }
43                let z = p[v];
44                p[v] = p[u];
45                v = z;
46            }
47        }
48    }
49
50    pub fn find(&mut self, mut u: usize) -> usize {
51        let mut p = self.parent[u];
52        while p != u {
53            self.parent[u] = self.parent[p];
54            u = p;
55            p = self.parent[p];
56        }
57
58        return u;
59    }
60
61    pub fn n_components(& self) -> usize {
62        return self.n_components;
63    }
64
65    pub fn n_nodes(&self) -> usize{
66        return self.parent.len();
67    }
68}