atcoder/lib/data_structure/
union_find.rs

1//! 素集合データ構造（Disjoint Set Union）です。
2
3/// `0..n` の要素を互いに素な集合へ分割して管理します。
4///
5/// 経路圧縮とサイズによる併合を使用します。各操作の償却計算量は `O(α(N))` です。
6/// 添字は 0-indexed で、範囲外の添字を渡すと panic します。
7///
8/// # Examples
9///
10/// ```
11/// use atcoder::data_structure::union_find::UnionFind;
12///
13/// let mut dsu = UnionFind::new(4);
14/// assert!(dsu.merge(0, 1));
15/// assert!(dsu.same(0, 1));
16/// assert_eq!(dsu.size(0), 2);
17/// assert_eq!(dsu.group_count(), 3);
18/// ```
19#[derive(Debug, Clone)]
20pub struct UnionFind {
21    parents: Vec<isize>,
22    group_count: usize,
23}
24
25impl UnionFind {
26    /// `0..n` の各要素を含む単一要素集合を作ります。計算量・メモリ量は `O(N)` です。
27    pub fn new(n: usize) -> Self {
28        Self {
29            parents: vec![-1; n],
30            group_count: n,
31        }
32    }
33
34    /// `x` の属する集合の代表元を返します。経路圧縮を行い、償却計算量は `O(α(N))` です。
35    /// `x >= n` なら panic します。
36    pub fn find(&mut self, x: usize) -> usize {
37        if self.parents[x] < 0 {
38            x
39        } else {
40            let root = self.find(self.parents[x] as usize);
41            self.parents[x] = root as isize;
42            root
43        }
44    }
45
46    /// `x` と `y` の集合を併合します。
47    ///
48    /// 新しく併合した場合は `true`、既に同じ集合なら `false` を返します。
49    /// 償却計算量は `O(α(N))` です。`x >= n` または `y >= n` なら panic します。
50    pub fn merge(&mut self, x: usize, y: usize) -> bool {
51        let mut x = self.find(x);
52        let mut y = self.find(y);
53        if x == y {
54            return false;
55        }
56        if self.parents[x] > self.parents[y] {
57            std::mem::swap(&mut x, &mut y);
58        }
59        self.parents[x] += self.parents[y];
60        self.parents[y] = x as isize;
61        self.group_count -= 1;
62        true
63    }
64
65    /// `x` と `y` が同じ集合なら `true` を返します。代表元の検索を行うため、
66    /// 償却計算量は `O(α(N))` です。`x >= n` または `y >= n` なら panic します。
67    pub fn same(&mut self, x: usize, y: usize) -> bool {
68        self.find(x) == self.find(y)
69    }
70
71    /// `x` の属する集合の要素数を返します。代表元の検索を行うため、
72    /// 償却計算量は `O(α(N))` です。`x >= n` なら panic します。
73    pub fn size(&mut self, x: usize) -> usize {
74        let root = self.find(x);
75        (-self.parents[root]) as usize
76    }
77
78    /// 現在の集合数を返します。計算量は `O(1)` です。
79    pub const fn group_count(&self) -> usize {
80        self.group_count
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use super::UnionFind;
87
88    #[test]
89    fn merges_components() {
90        let mut dsu = UnionFind::new(3);
91        assert!(dsu.merge(0, 2));
92        assert!(!dsu.merge(2, 0));
93        assert_eq!(dsu.size(2), 2);
94        assert_eq!(dsu.group_count(), 2);
95    }
96}
atcoder/lib/data_structure/union_find.rs

atcoder/lib/data_structure/
union_find.rs
