competitive/data_structure/
btreemap_ext.rs

1pub trait BTreeMapExt<K, V>
2where
3    K: Ord,
4{
5    fn first(&self) -> Option<(&K, &V)>;
6    fn last(&self) -> Option<(&K, &V)>;
7    fn get_next(&self, key: &K) -> Option<(&K, &V)>;
8    fn get_next_excluded(&self, key: &K) -> Option<(&K, &V)>;
9    fn get_next_back(&self, key: &K) -> Option<(&K, &V)>;
10    fn get_next_back_excluded(&self, key: &K) -> Option<(&K, &V)>;
11
12    fn first_mut(&mut self) -> Option<(&K, &mut V)>;
13    fn last_mut(&mut self) -> Option<(&K, &mut V)>;
14    fn get_next_mut(&mut self, key: &K) -> Option<(&K, &mut V)>;
15    fn get_next_excluded_mut(&mut self, key: &K) -> Option<(&K, &mut V)>;
16    fn get_next_back_mut(&mut self, key: &K) -> Option<(&K, &mut V)>;
17    fn get_next_back_excluded_mut(&mut self, key: &K) -> Option<(&K, &mut V)>;
18
19    fn pop_first(&mut self) -> Option<(K, V)>
20    where
21        K: Clone;
22    fn pop_last(&mut self) -> Option<(K, V)>
23    where
24        K: Clone;
25    fn pop_next(&mut self, key: &K) -> Option<(K, V)>
26    where
27        K: Clone;
28    fn pop_next_excluded(&mut self, key: &K) -> Option<(K, V)>
29    where
30        K: Clone;
31    fn pop_next_back(&mut self, key: &K) -> Option<(K, V)>
32    where
33        K: Clone;
34    fn pop_next_back_excluded(&mut self, key: &K) -> Option<(K, V)>
35    where
36        K: Clone;
37
38    fn pop_first_if<P>(&mut self, pred: P) -> Option<(K, V)>
39    where
40        K: Clone,
41        P: FnMut(&K, &V) -> bool;
42    fn pop_last_if<P>(&mut self, pred: P) -> Option<(K, V)>
43    where
44        K: Clone,
45        P: FnMut(&K, &V) -> bool;
46    fn pop_next_if<P>(&mut self, key: &K, pred: P) -> Option<(K, V)>
47    where
48        K: Clone,
49        P: FnMut(&K, &V) -> bool;
50    fn pop_next_excluded_if<P>(&mut self, key: &K, pred: P) -> Option<(K, V)>
51    where
52        K: Clone,
53        P: FnMut(&K, &V) -> bool;
54    fn pop_next_back_if<P>(&mut self, key: &K, pred: P) -> Option<(K, V)>
55    where
56        K: Clone,
57        P: FnMut(&K, &V) -> bool;
58    fn pop_next_back_excluded_if<P>(&mut self, key: &K, pred: P) -> Option<(K, V)>
59    where
60        K: Clone,
61        P: FnMut(&K, &V) -> bool;
62}
63impl<K, V> BTreeMapExt<K, V> for std::collections::BTreeMap<K, V>
64where
65    K: Ord,
66{
67    fn first(&self) -> Option<(&K, &V)> {
68        self.range(..).next()
69    }
70    fn last(&self) -> Option<(&K, &V)> {
71        self.range(..).next_back()
72    }
73    fn get_next(&self, key: &K) -> Option<(&K, &V)> {
74        self.range(key..).next()
75    }
76    fn get_next_excluded(&self, key: &K) -> Option<(&K, &V)> {
77        self.range((std::ops::Bound::Excluded(key), std::ops::Bound::Unbounded))
78            .next()
79    }
80    fn get_next_back(&self, key: &K) -> Option<(&K, &V)> {
81        self.range(..=key).next_back()
82    }
83    fn get_next_back_excluded(&self, key: &K) -> Option<(&K, &V)> {
84        self.range(..key).next_back()
85    }
86
87    fn first_mut(&mut self) -> Option<(&K, &mut V)> {
88        self.range_mut(..).next()
89    }
90    fn last_mut(&mut self) -> Option<(&K, &mut V)> {
91        self.range_mut(..).next_back()
92    }
93    fn get_next_mut(&mut self, key: &K) -> Option<(&K, &mut V)> {
94        self.range_mut(key..).next()
95    }
96    fn get_next_excluded_mut(&mut self, key: &K) -> Option<(&K, &mut V)> {
97        self.range_mut((std::ops::Bound::Excluded(key), std::ops::Bound::Unbounded))
98            .next()
99    }
100    fn get_next_back_mut(&mut self, key: &K) -> Option<(&K, &mut V)> {
101        self.range_mut(..=key).next_back()
102    }
103    fn get_next_back_excluded_mut(&mut self, key: &K) -> Option<(&K, &mut V)> {
104        self.range_mut(..key).next_back()
105    }
106
107    fn pop_first(&mut self) -> Option<(K, V)>
108    where
109        K: Clone,
110    {
111        self.pop_first_if(|_, _| true)
112    }
113    fn pop_last(&mut self) -> Option<(K, V)>
114    where
115        K: Clone,
116    {
117        self.pop_last_if(|_, _| true)
118    }
119    fn pop_next(&mut self, key: &K) -> Option<(K, V)>
120    where
121        K: Clone,
122    {
123        self.pop_next_if(key, |_, _| true)
124    }
125    fn pop_next_excluded(&mut self, key: &K) -> Option<(K, V)>
126    where
127        K: Clone,
128    {
129        self.pop_next_excluded_if(key, |_, _| true)
130    }
131    fn pop_next_back(&mut self, key: &K) -> Option<(K, V)>
132    where
133        K: Clone,
134    {
135        self.pop_next_back_if(key, |_, _| true)
136    }
137    fn pop_next_back_excluded(&mut self, key: &K) -> Option<(K, V)>
138    where
139        K: Clone,
140    {
141        self.pop_next_back_excluded_if(key, |_, _| true)
142    }
143
144    fn pop_first_if<P>(&mut self, mut pred: P) -> Option<(K, V)>
145    where
146        K: Clone,
147        P: FnMut(&K, &V) -> bool,
148    {
149        match self.first().filter(|(k, v)| pred(k, v)) {
150            Some((k, _)) => {
151                let k = k.clone();
152                let v = self.remove(&k).expect("This key must be exists.");
153                Some((k, v))
154            }
155            None => None,
156        }
157    }
158    fn pop_last_if<P>(&mut self, mut pred: P) -> Option<(K, V)>
159    where
160        K: Clone,
161        P: FnMut(&K, &V) -> bool,
162    {
163        match self.last().filter(|(k, v)| pred(k, v)) {
164            Some((k, _)) => {
165                let k = k.clone();
166                let v = self.remove(&k).expect("This key must be exists.");
167                Some((k, v))
168            }
169            None => None,
170        }
171    }
172    fn pop_next_if<P>(&mut self, key: &K, mut pred: P) -> Option<(K, V)>
173    where
174        K: Clone,
175        P: FnMut(&K, &V) -> bool,
176    {
177        match self.get_next(key).filter(|(k, v)| pred(k, v)) {
178            Some((k, _)) => {
179                let k = k.clone();
180                let v = self.remove(&k).expect("This key must be exists.");
181                Some((k, v))
182            }
183            None => None,
184        }
185    }
186    fn pop_next_excluded_if<P>(&mut self, key: &K, mut pred: P) -> Option<(K, V)>
187    where
188        K: Clone,
189        P: FnMut(&K, &V) -> bool,
190    {
191        match self.get_next_excluded(key).filter(|(k, v)| pred(k, v)) {
192            Some((k, _)) => {
193                let k = k.clone();
194                let v = self.remove(&k).expect("This key must be exists.");
195                Some((k, v))
196            }
197            None => None,
198        }
199    }
200    fn pop_next_back_if<P>(&mut self, key: &K, mut pred: P) -> Option<(K, V)>
201    where
202        K: Clone,
203        P: FnMut(&K, &V) -> bool,
204    {
205        match self.get_next_back(key).filter(|(k, v)| pred(k, v)) {
206            Some((k, _)) => {
207                let k = k.clone();
208                let v = self.remove(&k).expect("This key must be exists.");
209                Some((k, v))
210            }
211            None => None,
212        }
213    }
214    fn pop_next_back_excluded_if<P>(&mut self, key: &K, mut pred: P) -> Option<(K, V)>
215    where
216        K: Clone,
217        P: FnMut(&K, &V) -> bool,
218    {
219        match self.get_next_back_excluded(key).filter(|(k, v)| pred(k, v)) {
220            Some((k, _)) => {
221                let k = k.clone();
222                let v = self.remove(&k).expect("This key must be exists.");
223                Some((k, v))
224            }
225            None => None,
226        }
227    }
228}
229
230pub trait BTreeSetExt<T>
231where
232    T: Ord,
233{
234    fn first(&self) -> Option<&T>;
235    fn last(&self) -> Option<&T>;
236    fn get_next(&self, key: &T) -> Option<&T>;
237    fn get_next_excluded(&self, key: &T) -> Option<&T>;
238    fn get_next_back(&self, key: &T) -> Option<&T>;
239    fn get_next_back_excluded(&self, key: &T) -> Option<&T>;
240
241    fn pop_first(&mut self) -> Option<T>
242    where
243        T: Clone;
244    fn pop_last(&mut self) -> Option<T>
245    where
246        T: Clone;
247    fn pop_next(&mut self, key: &T) -> Option<T>
248    where
249        T: Clone;
250    fn pop_next_excluded(&mut self, key: &T) -> Option<T>
251    where
252        T: Clone;
253    fn pop_next_back(&mut self, key: &T) -> Option<T>
254    where
255        T: Clone;
256    fn pop_next_back_excluded(&mut self, key: &T) -> Option<T>
257    where
258        T: Clone;
259
260    fn pop_first_if<P>(&mut self, pred: P) -> Option<T>
261    where
262        T: Clone,
263        P: FnMut(&T) -> bool;
264    fn pop_last_if<P>(&mut self, pred: P) -> Option<T>
265    where
266        T: Clone,
267        P: FnMut(&T) -> bool;
268    fn pop_next_if<P>(&mut self, key: &T, pred: P) -> Option<T>
269    where
270        T: Clone,
271        P: FnMut(&T) -> bool;
272    fn pop_next_excluded_if<P>(&mut self, key: &T, pred: P) -> Option<T>
273    where
274        T: Clone,
275        P: FnMut(&T) -> bool;
276    fn pop_next_back_if<P>(&mut self, key: &T, pred: P) -> Option<T>
277    where
278        T: Clone,
279        P: FnMut(&T) -> bool;
280    fn pop_next_back_excluded_if<P>(&mut self, key: &T, pred: P) -> Option<T>
281    where
282        T: Clone,
283        P: FnMut(&T) -> bool;
284}
285impl<T> BTreeSetExt<T> for std::collections::BTreeSet<T>
286where
287    T: Ord,
288{
289    fn first(&self) -> Option<&T> {
290        self.range(..).next()
291    }
292    fn last(&self) -> Option<&T> {
293        self.range(..).next_back()
294    }
295    fn get_next(&self, key: &T) -> Option<&T> {
296        self.range(key..).next()
297    }
298    fn get_next_excluded(&self, key: &T) -> Option<&T> {
299        self.range((std::ops::Bound::Excluded(key), std::ops::Bound::Unbounded))
300            .next()
301    }
302    fn get_next_back(&self, key: &T) -> Option<&T> {
303        self.range(..=key).next_back()
304    }
305    fn get_next_back_excluded(&self, key: &T) -> Option<&T> {
306        self.range(..key).next_back()
307    }
308
309    fn pop_first(&mut self) -> Option<T>
310    where
311        T: Clone,
312    {
313        self.pop_first_if(|_| true)
314    }
315    fn pop_last(&mut self) -> Option<T>
316    where
317        T: Clone,
318    {
319        self.pop_last_if(|_| true)
320    }
321    fn pop_next(&mut self, key: &T) -> Option<T>
322    where
323        T: Clone,
324    {
325        self.pop_next_if(key, |_| true)
326    }
327    fn pop_next_excluded(&mut self, key: &T) -> Option<T>
328    where
329        T: Clone,
330    {
331        self.pop_next_excluded_if(key, |_| true)
332    }
333    fn pop_next_back(&mut self, key: &T) -> Option<T>
334    where
335        T: Clone,
336    {
337        self.pop_next_back_if(key, |_| true)
338    }
339    fn pop_next_back_excluded(&mut self, key: &T) -> Option<T>
340    where
341        T: Clone,
342    {
343        self.pop_next_back_excluded_if(key, |_| true)
344    }
345
346    fn pop_first_if<P>(&mut self, mut pred: P) -> Option<T>
347    where
348        T: Clone,
349        P: FnMut(&T) -> bool,
350    {
351        match <Self as BTreeSetExt<T>>::first(self).filter(|k| pred(k)) {
352            Some(k) => {
353                let k = k.clone();
354                self.remove(&k);
355                Some(k)
356            }
357            None => None,
358        }
359    }
360    fn pop_last_if<P>(&mut self, mut pred: P) -> Option<T>
361    where
362        T: Clone,
363        P: FnMut(&T) -> bool,
364    {
365        match <Self as BTreeSetExt<T>>::last(self).filter(|k| pred(k)) {
366            Some(k) => {
367                let k = k.clone();
368                self.remove(&k);
369                Some(k)
370            }
371            None => None,
372        }
373    }
374    fn pop_next_if<P>(&mut self, key: &T, mut pred: P) -> Option<T>
375    where
376        T: Clone,
377        P: FnMut(&T) -> bool,
378    {
379        match self.get_next(key).filter(|k| pred(k)) {
380            Some(k) => {
381                let k = k.clone();
382                self.remove(&k);
383                Some(k)
384            }
385            None => None,
386        }
387    }
388    fn pop_next_excluded_if<P>(&mut self, key: &T, mut pred: P) -> Option<T>
389    where
390        T: Clone,
391        P: FnMut(&T) -> bool,
392    {
393        match self.get_next_excluded(key).filter(|k| pred(k)) {
394            Some(k) => {
395                let k = k.clone();
396                self.remove(&k);
397                Some(k)
398            }
399            None => None,
400        }
401    }
402    fn pop_next_back_if<P>(&mut self, key: &T, mut pred: P) -> Option<T>
403    where
404        T: Clone,
405        P: FnMut(&T) -> bool,
406    {
407        match self.get_next_back(key).filter(|k| pred(k)) {
408            Some(k) => {
409                let k = k.clone();
410                self.remove(&k);
411                Some(k)
412            }
413            None => None,
414        }
415    }
416    fn pop_next_back_excluded_if<P>(&mut self, key: &T, mut pred: P) -> Option<T>
417    where
418        T: Clone,
419        P: FnMut(&T) -> bool,
420    {
421        match self.get_next_back_excluded(key).filter(|k| pred(k)) {
422            Some(k) => {
423                let k = k.clone();
424                self.remove(&k);
425                Some(k)
426            }
427            None => None,
428        }
429    }
430}