1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
// src/unionfind.rs
use pyo3::prelude::*;

/// The Union-Find structure is used to efficiently manage and merge disjoint sets.
/// This module provides an implementation of the Union-Find (disjoint set) data structure.
#[pyclass]
pub struct UnionFind {
    /// The parent array, where parent[i] points to the parent of node i. If parent[i] == i, then i is a root.
    parent: Vec<usize>,
    /// The rank array, used to keep the tree flat by storing the depth of the tree for each node.

    rank: Vec<usize>,
}

#[pymethods]
impl UnionFind {

    /// Creates a new Union-Find data structure with `n` elements, each initially in its own set.
    ///
    /// # Arguments
    ///
    /// * `n` - The number of elements in the Union-Find data structure.
    ///
    /// # Example
    ///
    /// ```
    /// let uf = UnionFind::new(10);
    /// ```
    #[new]
    pub fn new(n: usize) -> Self {
        let mut parent = vec![0; n];
        let mut rank = vec![0; n];
        for i in 0..n {
            parent[i] = i;
        }
        UnionFind { parent, rank }
    }
    /// Finds the representative (root) of the set containing `x` using path compression.
    ///
    /// # Arguments
    ///
    /// * `x` - The element for which to find the set representative.
    ///
    /// # Returns
    ///
    /// The representative of the set containing `x`.
    ///
    /// # Example
    ///
    /// ```
    /// let root = uf.find(3);
    /// ```
    pub fn find(&mut self, x: usize) -> usize {
        if self.parent[x] != x {
            self.parent[x] = self.find(self.parent[x]);
        }
        self.parent[x]
    }
    /// Unites the sets containing `x` and `y`. Uses union by rank to keep the tree flat.
    ///
    /// # Arguments
    ///
    /// * `x` - An element in the first set.
    /// * `y` - An element in the second set.
    ///
    /// # Example
    ///
    /// ```
    /// uf.union(3, 4);
    /// ```
    pub fn union(&mut self, x: usize, y: usize) {
        let root_x = self.find(x);
        let root_y = self.find(y);

        if root_x == root_y {
            return;
        }

        if self.rank[root_x] < self.rank[root_y] {
            self.parent[root_x] = root_y;
        } else if self.rank[root_x] > self.rank[root_y] {
            self.parent[root_y] = root_x;
        } else {
            self.parent[root_y] = root_x;
            self.rank[root_x] += 1;
        }
    }
    /// Checks if `x` and `y` are in the same set.
    ///
    /// # Arguments
    ///
    /// * `x` - An element in the first set.
    /// * `y` - An element in the second set.
    ///
    /// # Returns
    ///
    /// `true` if `x` and `y` are in the same set, `false` otherwise.
    ///
    /// # Example
    ///
    /// ```
    /// let same = uf.same_set(3, 4);
    /// ```
    pub fn same_set(&mut self, x: usize, y: usize) -> bool {
        self.find(x) == self.find(y)
    }
}