haar_lib/ds/
unionfind.rs

1//! 素集合データ構造
2//!
3//! # Problems
4//! - <https://atcoder.jp/contests/abc372/tasks/abc372_e>
5use std::cell::Cell;
6
7/// 素集合の統合と所属の判定ができるデータ構造。
8pub struct UnionFind<'a, T = ()> {
9    n: usize,
10    count: usize,
11    parent: Vec<Cell<usize>>,
12    depth: Vec<usize>,
13    size: Vec<usize>,
14    values: Option<Vec<Option<T>>>,
15    merge: Option<Box<dyn 'a + Fn(T, T) -> T>>,
16}
17
18impl UnionFind<'_, ()> {
19    /// 大きさ`1`の集合を`n`個用意する。
20    pub fn new(n: usize) -> Self {
21        UnionFind {
22            n,
23            count: n,
24            parent: (0..n).map(Cell::new).collect(),
25            depth: vec![1; n],
26            size: vec![1; n],
27            values: None,
28            merge: None,
29        }
30    }
31}
32
33impl<'a, T> UnionFind<'a, T> {
34    /// 大きさ`1`の集合を`|values|`個用意する。このとき、各集合`i`に`value[i]`を割り当てる。
35    ///
36    /// `merge`は、集合を統合する際に、新しい集合に割り当てる値を返す。
37    pub fn with_values(values: Vec<T>, merge: Box<impl 'a + Fn(T, T) -> T>) -> Self {
38        let n = values.len();
39        UnionFind {
40            n,
41            count: n,
42            parent: (0..n).map(Cell::new).collect(),
43            depth: vec![1; n],
44            size: vec![1; n],
45            values: Some(values.into_iter().map(Option::Some).collect()),
46            merge: Some(Box::new(merge)),
47        }
48    }
49
50    /// `i`の属する集合の根を返す。
51    pub fn root_of(&self, i: usize) -> usize {
52        if self.parent[i].get() == i {
53            return i;
54        }
55        let p = self.parent[i].get();
56        self.parent[i].set(self.root_of(p));
57        self.parent[i].get()
58    }
59
60    /// `i`と`j`が同じ集合に属するならば`true`を返す。
61    pub fn is_same(&self, i: usize, j: usize) -> bool {
62        self.root_of(i) == self.root_of(j)
63    }
64
65    /// `i`の属する集合と`j`の属する集合を統合する。
66    pub fn merge(&mut self, i: usize, j: usize) -> usize {
67        let i = self.root_of(i);
68        let j = self.root_of(j);
69
70        if i == j {
71            return i;
72        }
73
74        let (p, c) = if self.depth[i] < self.depth[j] {
75            (j, i)
76        } else {
77            (i, j)
78        };
79
80        self.count -= 1;
81
82        self.parent[c].set(p);
83        self.size[p] += self.size[c];
84        if self.depth[p] == self.depth[c] {
85            self.depth[p] += 1;
86        }
87
88        if let Some(f) = self.merge.as_ref() {
89            let t = f(
90                self.values.as_mut().unwrap()[p].take().unwrap(),
91                self.values.as_mut().unwrap()[c].take().unwrap(),
92            );
93            self.values.as_mut().unwrap()[p] = Some(t);
94        }
95
96        p
97    }
98
99    /// `i`の属する集合の大きさを返す。
100    pub fn size_of(&self, i: usize) -> usize {
101        let i = self.root_of(i);
102        self.size[i]
103    }
104
105    /// 素集合の個数を返す。
106    pub fn count_groups(&self) -> usize {
107        self.count
108    }
109
110    /// `i`の属する集合のもつ値を返す。
111    pub fn value_of(&self, i: usize) -> Option<&T> {
112        let i = self.root_of(i);
113        self.values.as_ref()?[i].as_ref()
114    }
115
116    /// `i`の属する集合のもつ値への可変参照を返す。
117    pub fn value_mut_of(&mut self, i: usize) -> Option<&mut T> {
118        let i = self.root_of(i);
119        self.values.as_mut()?[i].as_mut()
120    }
121
122    /// 素集合をすべて列挙する。
123    pub fn get_groups(&self) -> Vec<Vec<usize>> {
124        let mut ret = vec![vec![]; self.n];
125
126        for i in 0..self.n {
127            ret[self.root_of(i)].push(i);
128        }
129
130        ret.into_iter().filter(|x| !x.is_empty()).collect()
131    }
132}
133
134#[cfg(test)]
135mod tests {
136    use super::*;
137    use crate::btreeset;
138    use rand::Rng;
139    use std::collections::BTreeSet;
140    use std::iter::FromIterator;
141
142    #[test]
143    fn test() {
144        let n = 100;
145        let q = 50;
146        let mut rng = rand::thread_rng();
147
148        let mut uf = UnionFind::new(n);
149        let mut a = (0..n).map(|i| btreeset![i]).collect::<BTreeSet<_>>();
150
151        for _ in 0..q {
152            let i = rng.gen_range(0..n);
153            let j = rng.gen_range(0..n);
154
155            uf.merge(i, j);
156
157            let mut ai = a.iter().find(|s| s.contains(&i)).unwrap().clone();
158            let aj = a.iter().find(|s| s.contains(&j)).unwrap().clone();
159
160            if ai != aj {
161                a.remove(&ai);
162                a.remove(&aj);
163                ai.extend(aj);
164                a.insert(ai);
165            }
166        }
167
168        for _ in 0..q {
169            let i = rng.gen_range(0..n);
170            let j = rng.gen_range(0..n);
171
172            let ai = a.iter().find(|s| s.contains(&i)).unwrap();
173
174            assert_eq!(uf.is_same(i, j), ai.contains(&j));
175        }
176
177        assert_eq!(
178            BTreeSet::from_iter(
179                uf.get_groups()
180                    .into_iter()
181                    .map(|s| BTreeSet::from_iter(s.into_iter()))
182            ),
183            a
184        );
185    }
186}
haar_lib/ds/unionfind.rs

haar_lib/ds/
unionfind.rs
