competitive/data_structure/
sparse_set.rs

1#[derive(Debug, Clone)]
2pub struct SparseSet<Usize, const FIXED: bool = true> {
3    data: Vec<Usize>,
4    index: Vec<Usize>,
5}
6
7macro_rules! impl_sparse_set {
8    ($($ty:ty)*) => {
9        $(
10            impl<const FIXED: bool> SparseSet<$ty, FIXED> {
11                pub fn new(n: usize) -> Self {
12                    Self {
13                        data: Vec::with_capacity(n),
14                        index: vec![!0; n],
15                    }
16                }
17
18                pub unsafe fn insert_unchecked(&mut self, x: $ty) {
19                    if !FIXED && x as usize >= self.index.len() {
20                        self.index.resize_with(x as usize + 1, Default::default);
21                    }
22                    self.index[x as usize] = self.data.len() as _;
23                    self.data.push(x);
24                }
25
26                pub fn insert(&mut self, x: $ty) -> bool {
27                    if self.contains(x) {
28                        return false;
29                    }
30                    unsafe { self.insert_unchecked(x) };
31                    true
32                }
33
34                pub unsafe fn remove_unchecked(&mut self, x: $ty) {
35                    let n = self.data.len();
36                    let k = std::mem::replace(&mut self.index[x as usize], !0);
37                    if k as usize != n - 1 {
38                        self.data.swap(k as usize, n - 1);
39                        self.index[self.data[k as usize] as usize] = k;
40                    }
41                    self.data.pop();
42                }
43
44                pub fn remove(&mut self, x: $ty) -> bool {
45                    if !self.contains(x) {
46                        return false;
47                    }
48                    unsafe { self.remove_unchecked(x) };
49                    true
50                }
51
52                pub fn len(&self) -> usize {
53                    self.data.len()
54                }
55
56                pub fn is_empty(&self) -> bool {
57                    self.len() == 0
58                }
59
60                pub fn contains(&self, x: $ty) -> bool {
61                    self.index[x as usize] != !0
62                }
63
64                pub fn iter(&self) -> std::slice::Iter<'_, $ty> {
65                    self.data.iter()
66                }
67            }
68        )*
69    };
70}
71impl_sparse_set!(u8 u16 u32 u64 usize);