cp_library_rs/data_structure/
mex_set.rs

1//! Mexを管理するデータ構造
2
3use std::{
4    collections::BTreeSet,
5    ops::{Bound::*, RangeBounds},
6};
7
8/// 集合とそのmexを管理する
9#[derive(Debug)]
10pub struct MexSet {
11    pub ranges: BTreeSet<(isize, isize)>,
12}
13
14impl Default for MexSet {
15    fn default() -> Self {
16        Self::new()
17    }
18}
19
20impl MexSet {
21    const INF: isize = isize::MIN;
22    const SUP: isize = isize::MAX;
23
24    /// MexSetを初期化する
25    pub fn new() -> Self {
26        let ranges = [(Self::INF, Self::INF), (Self::SUP, Self::SUP)]
27            .into_iter()
28            .collect();
29        Self { ranges }
30    }
31
32    /// 集合に要素`x`を追加する
33    /// ### 戻り値
34    /// - `true`: `x`が追加された場合
35    /// - `false`: `x`がすでに存在していた場合
36    pub fn insert(&mut self, x: isize) -> bool {
37        // 範囲が`(INF, SUP)` である場合 → 追加されない
38        if self.ranges.len() == 1 {
39            return false;
40        }
41        let &(ll, l) = self.ranges.range(..(x + 1, x + 1)).next_back().unwrap();
42        let &(r, rr) = self.ranges.range((x + 1, x + 1)..).next().unwrap();
43        if x <= l {
44            return false;
45        }
46        match (l == x - 1, x + 1 == r) {
47            (false, false) => {
48                self.ranges.insert((x, x));
49            }
50            (false, true) => {
51                self.ranges.remove(&(r, rr));
52                self.ranges.insert((x, rr));
53            }
54            (true, false) => {
55                self.ranges.remove(&(ll, l));
56                self.ranges.insert((ll, x));
57            }
58            (true, true) => {
59                self.ranges.remove(&(ll, l));
60                self.ranges.remove(&(r, rr));
61                self.ranges.insert((ll, rr));
62            }
63        }
64        true
65    }
66
67    /// 集合から要素`x`を削除する
68    /// ### 戻り値
69    /// - `true`: `x`が削除された場合
70    /// - `false`: `x`がすでに存在していなかった場合
71    pub fn delete(&mut self, x: isize) -> bool {
72        let &(ll, l) = self.ranges.range(..(x + 1, x + 1)).next_back().unwrap();
73        if l < x {
74            return false;
75        }
76        self.ranges.remove(&(ll, l));
77        match (ll == x, x == l) {
78            (false, false) => {
79                self.ranges.insert((ll, x - 1));
80                self.ranges.insert((x + 1, l));
81            }
82            (false, true) => {
83                self.ranges.insert((ll, x - 1));
84            }
85            (true, false) => {
86                self.ranges.insert((x + 1, l));
87            }
88            (true, true) => (),
89        }
90        true
91    }
92
93    #[inline]
94    fn parse_range<R: RangeBounds<isize>>(range: R) -> (isize, isize) {
95        let start = match range.start_bound() {
96            Unbounded => Self::INF,
97            Included(&s) => s,
98            Excluded(&s) => s + 1,
99        };
100        let end = match range.end_bound() {
101            Unbounded => Self::SUP,
102            Included(&e) => e,
103            Excluded(&e) => e - 1,
104        };
105        (start, end)
106    }
107
108    /// 集合に区間を追加する
109    /// - 計算量: O(log(n)) (amotized)
110    pub fn insert_range<R: RangeBounds<isize>>(&mut self, range: R) -> bool {
111        let (start, end) = Self::parse_range(range);
112        if start > end {
113            return false;
114        }
115        // 範囲が`(INF, SUP)` である場合 → 追加されない
116        if self.ranges.len() == 1 {
117            return false;
118        }
119        // 全範囲を追加する場合
120        if (start, end) == (Self::INF, Self::SUP) {
121            let len = self.ranges.len();
122            self.ranges.clear();
123            self.ranges.insert((Self::INF, Self::SUP));
124            return len != self.ranges.len();
125        }
126        if start == end {
127            return self.insert(start);
128        }
129        while let Some(&(l, r)) = self
130            .ranges
131            .range((Excluded((start, start)), Excluded((end, end))))
132            .next()
133        {
134            if r < end {
135                self.ranges.remove(&(l, r));
136            } else {
137                break;
138            }
139        }
140        let &(ll, l) = self
141            .ranges
142            .range(..(start + 1, start + 1))
143            .next_back()
144            .unwrap();
145        let &(r, rr) = self.ranges.range((end, end)..).next().unwrap();
146        if ll <= start && end <= l {
147            return false;
148        }
149        match (start <= l + 1, r - 1 <= end) {
150            (false, false) => {
151                self.ranges.insert((start, end));
152            }
153            (false, true) => {
154                self.ranges.remove(&(r, rr));
155                self.ranges.insert((start, rr));
156            }
157            (true, false) => {
158                self.ranges.remove(&(ll, l));
159                self.ranges.insert((ll, end));
160            }
161            (true, true) => {
162                self.ranges.remove(&(ll, l));
163                self.ranges.remove(&(r, rr));
164                self.ranges.insert((ll, rr));
165            }
166        }
167        true
168    }
169
170    /// 集合から区間を削除する
171    /// - 計算量: O(log(n)) (amotized)
172    pub fn delete_range<R: RangeBounds<isize>>(&mut self, range: R) -> bool {
173        let (start, end) = Self::parse_range(range);
174        if start > end {
175            return false;
176        }
177        // 全範囲の場合
178        if (start, end) == (Self::INF, Self::SUP) {
179            let len = self.ranges.len();
180            self.ranges.clear();
181            self.ranges.insert((Self::INF, Self::INF));
182            self.ranges.insert((Self::SUP, Self::SUP));
183            return len != self.ranges.len();
184        }
185        if start == end {
186            return self.delete(start);
187        }
188        while let Some(&(l, r)) = self
189            .ranges
190            .range((Excluded((start, start)), Excluded((end, end))))
191            .next()
192        {
193            if r < end {
194                self.ranges.remove(&(l, r));
195            } else {
196                break;
197            }
198        }
199        let &(ll, l) = self
200            .ranges
201            .range((Unbounded, Included((start, start))))
202            .next_back()
203            .unwrap();
204        let &(r, rr) = self
205            .ranges
206            .range((Unbounded, Included((end, end))))
207            .next_back()
208            .unwrap();
209        if l < start && end < r {
210            return false;
211        }
212        if start <= l {
213            self.ranges.remove(&(ll, l));
214            match (ll < start, end < l) {
215                (false, false) => {}
216                (false, true) => {
217                    self.ranges.insert((end.saturating_add(1), l));
218                }
219                (true, false) => {
220                    self.ranges.insert((ll, start.saturating_sub(1)));
221                }
222                (true, true) => {
223                    self.ranges.insert((ll, start.saturating_sub(1)));
224                    self.ranges.insert((end.saturating_add(1), l));
225                }
226            }
227        }
228        if r <= end && end <= rr {
229            self.ranges.remove(&(r, rr));
230            self.ranges.insert((end.saturating_add(1), rr));
231        }
232        true
233    }
234
235    /// **集合に含まれない**`x`以上で最小の整数を調べる
236    pub fn mex(&self, x: isize) -> isize {
237        let &(ll, l) = self.ranges.range(..(x + 1, x + 1)).next_back().unwrap();
238        if ll <= x && x <= l {
239            l + 1
240        } else {
241            x
242        }
243    }
244
245    /// **集合に含まれない**`x`以下で最大の整数を調べる
246    pub fn mex_rev(&self, x: isize) -> isize {
247        let &(l, r) = self.ranges.range(..(x, isize::MAX)).next_back().unwrap();
248        if l <= x && x <= r {
249            l - 1
250        } else {
251            x
252        }
253    }
254}
cp_library_rs/data_structure/mex_set.rs

cp_library_rs/data_structure/
mex_set.rs
