haar_lib/ds/
integer_set.rs

1//! Mexを求められるデータ構造
2//!
3//! # Problems
4//! - <https://atcoder.jp/contests/hhkb2020/tasks/hhkb2020_c>
5//!
6//! # References
7//! - <https://rsk0315.hatenablog.com/entry/2020/10/11/125049>
8
9use std::collections::BTreeMap;
10
11/// 整数の追加・削除とMexを求められるデータ構造
12#[derive(Clone, Debug, Default)]
13pub struct IntegerSet {
14    data: BTreeMap<i64, i64>,
15}
16
17impl IntegerSet {
18    /// 空の`IntegerSet`を生成
19    pub fn new() -> Self {
20        Self {
21            ..Default::default()
22        }
23    }
24
25    /// `x`を含む半開区間を返す
26    ///
27    /// **Time complexity** $O(\log n)$
28    pub fn interval(&self, x: i64) -> Option<(i64, i64)> {
29        if let Some((&k, &v)) = self.data.range(..=x).next_back() {
30            if k <= x && x < v {
31                return Some((k, v));
32            }
33        }
34        None
35    }
36
37    /// `x`を含むかを判定
38    ///
39    /// **Time complexity** $O(\log n)$
40    pub fn contains(&self, x: i64) -> bool {
41        self.interval(x).is_some()
42    }
43
44    /// `x`を追加する
45    ///
46    /// **Time complexity** $O(\log n)$
47    pub fn insert(&mut self, x: i64) {
48        if let Some((&k, &v)) = self.data.range(..=x).next_back() {
49            if k <= x && x < v {
50                return;
51            } else if x == v {
52                if self.data.contains_key(&(x + 1)) {
53                    let v = self.data.remove(&(x + 1)).unwrap();
54                    self.data.insert(k, v);
55                } else {
56                    self.data.insert(k, x + 1);
57                }
58                return;
59            }
60        }
61
62        if self.data.contains_key(&(x + 1)) {
63            let v = self.data.remove(&(x + 1)).unwrap();
64            self.data.insert(x, v);
65        } else {
66            self.data.insert(x, x + 1);
67        }
68    }
69
70    /// `x`を削除する
71    ///
72    /// **Time complexity** $O(\log n)$
73    pub fn remove(&mut self, x: i64) {
74        if let Some((k, v)) = self.interval(x) {
75            self.data.remove(&k);
76
77            if k != x {
78                self.data.insert(k, x);
79            }
80            if x + 1 != v {
81                self.data.insert(x + 1, v);
82            }
83        }
84    }
85
86    /// `self`に含まれない数のうち`x`以上で最小のもの
87    ///
88    /// **Time complexity** $O(\log n)$
89    pub fn mex(&self, x: i64) -> i64 {
90        self.interval(x).map_or(x, |(_, v)| v)
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use std::collections::BTreeSet;
97
98    use super::*;
99    use rand::Rng;
100
101    #[test]
102    fn test() {
103        let mut rng = rand::thread_rng();
104
105        let n = 100;
106        let t = 1000;
107        let l = 200;
108
109        let mut s = IntegerSet::new();
110        let mut a = BTreeSet::new();
111
112        for _ in 0..n {
113            let x = rng.gen_range(-l..=l);
114            s.insert(x);
115            a.insert(x);
116        }
117
118        for _ in 0..t {
119            let x = rng.gen_range(-l..=l);
120            s.insert(x);
121            a.insert(x);
122
123            let x = rng.gen_range(-l..=l);
124            s.remove(x);
125            a.remove(&x);
126
127            let x = rng.gen_range(-l..=l);
128
129            let mut mex = x;
130            while a.contains(&mex) {
131                mex += 1;
132            }
133
134            assert_eq!(s.mex(x), mex);
135        }
136    }
137}
haar_lib/ds/integer_set.rs

haar_lib/ds/
integer_set.rs
