competitive/data_structure/
container.rs

1use std::{
2    cmp::Ord,
3    collections::{BTreeMap, HashMap, btree_map, hash_map},
4    hash::Hash,
5    iter::FusedIterator,
6    marker::PhantomData,
7};
8
9pub trait ContainerFactory {
10    type Container: Container;
11
12    fn create_container(&self) -> Self::Container;
13}
14
15pub trait Container {
16    type Key;
17    type Value;
18    type Entry<'a>: ContainerEntry<'a, Key = Self::Key, Value = Self::Value>
19    where
20        Self: 'a,
21        Self::Key: 'a,
22        Self::Value: 'a;
23    type Iter<'a>: Iterator<Item = (&'a Self::Key, &'a Self::Value)>
24    where
25        Self: 'a,
26        Self::Key: 'a,
27        Self::Value: 'a;
28    type Drain<'a>: Iterator<Item = (Self::Key, Self::Value)>
29    where
30        Self: 'a,
31        Self::Key: 'a,
32        Self::Value: 'a;
33
34    fn get(&self, key: &Self::Key) -> Option<&Self::Value>;
35    fn get_mut(&mut self, key: &Self::Key) -> Option<&mut Self::Value>;
36    fn insert(&mut self, key: Self::Key, value: Self::Value) -> Option<Self::Value>;
37    fn remove(&mut self, key: &Self::Key) -> Option<Self::Value>;
38    fn entry(&mut self, key: Self::Key) -> Self::Entry<'_>;
39    fn iter(&self) -> Self::Iter<'_>;
40    fn drain(&mut self) -> Self::Drain<'_>;
41}
42
43pub trait ContainerEntry<'a> {
44    type Key: 'a;
45    type Value: 'a;
46
47    fn or_default(self) -> &'a mut Self::Value
48    where
49        Self::Value: Default;
50    fn or_insert(self, default: Self::Value) -> &'a mut Self::Value;
51    fn or_insert_with<F>(self, default: F) -> &'a mut Self::Value
52    where
53        F: FnOnce() -> Self::Value;
54    fn or_insert_with_key<F>(self, default: F) -> &'a mut Self::Value
55    where
56        F: FnOnce(&Self::Key) -> Self::Value;
57    fn key(&self) -> &Self::Key;
58    fn and_modify<F>(self, f: F) -> Self
59    where
60        F: FnOnce(&mut Self::Value);
61}
62
63impl<F> ContainerFactory for &F
64where
65    F: ContainerFactory,
66{
67    type Container = F::Container;
68
69    fn create_container(&self) -> Self::Container {
70        F::create_container(self)
71    }
72}
73
74#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
75pub struct BTreeMapFactory<K, V> {
76    _marker: PhantomData<fn() -> (K, V)>,
77}
78
79impl<K, V> Default for BTreeMapFactory<K, V> {
80    fn default() -> Self {
81        Self {
82            _marker: Default::default(),
83        }
84    }
85}
86
87#[derive(Debug)]
88pub struct BTreeMapDrain<'a, K, V> {
89    inner: &'a mut BTreeMap<K, V>,
90}
91
92impl<K, V> ContainerFactory for BTreeMapFactory<K, V>
93where
94    K: Ord,
95{
96    type Container = BTreeMap<K, V>;
97
98    fn create_container(&self) -> Self::Container {
99        BTreeMap::new()
100    }
101}
102
103impl<K, V> Container for BTreeMap<K, V>
104where
105    K: Ord,
106{
107    type Key = K;
108    type Value = V;
109    type Entry<'a>
110        = btree_map::Entry<'a, K, V>
111    where
112        K: 'a,
113        V: 'a;
114    type Iter<'a>
115        = btree_map::Iter<'a, K, V>
116    where
117        K: 'a,
118        V: 'a;
119    type Drain<'a>
120        = BTreeMapDrain<'a, K, V>
121    where
122        K: 'a,
123        V: 'a;
124
125    fn get(&self, key: &Self::Key) -> Option<&Self::Value> {
126        self.get(key)
127    }
128
129    fn get_mut(&mut self, key: &Self::Key) -> Option<&mut Self::Value> {
130        self.get_mut(key)
131    }
132
133    fn insert(&mut self, key: Self::Key, value: Self::Value) -> Option<Self::Value> {
134        self.insert(key, value)
135    }
136
137    fn remove(&mut self, key: &Self::Key) -> Option<Self::Value> {
138        self.remove(key)
139    }
140
141    fn entry(&mut self, key: Self::Key) -> Self::Entry<'_> {
142        self.entry(key)
143    }
144
145    fn iter(&self) -> Self::Iter<'_> {
146        self.iter()
147    }
148
149    fn drain(&mut self) -> Self::Drain<'_> {
150        BTreeMapDrain { inner: self }
151    }
152}
153
154impl<K, V> Iterator for BTreeMapDrain<'_, K, V>
155where
156    K: Ord,
157{
158    type Item = (K, V);
159
160    fn next(&mut self) -> Option<Self::Item> {
161        self.inner.pop_first()
162    }
163}
164
165impl<K, V> DoubleEndedIterator for BTreeMapDrain<'_, K, V>
166where
167    K: Ord,
168{
169    fn next_back(&mut self) -> Option<Self::Item> {
170        self.inner.pop_last()
171    }
172}
173
174impl<K, V> ExactSizeIterator for BTreeMapDrain<'_, K, V>
175where
176    K: Ord,
177{
178    fn len(&self) -> usize {
179        self.inner.len()
180    }
181}
182
183impl<K, V> FusedIterator for BTreeMapDrain<'_, K, V> where K: Ord {}
184
185impl<'a, K, V> ContainerEntry<'a> for btree_map::Entry<'a, K, V>
186where
187    K: 'a + Ord,
188    V: 'a,
189{
190    type Key = K;
191    type Value = V;
192
193    fn or_default(self) -> &'a mut Self::Value
194    where
195        Self::Value: Default,
196    {
197        self.or_default()
198    }
199
200    fn or_insert(self, default: Self::Value) -> &'a mut Self::Value {
201        self.or_insert(default)
202    }
203
204    fn or_insert_with<F>(self, default: F) -> &'a mut Self::Value
205    where
206        F: FnOnce() -> Self::Value,
207    {
208        self.or_insert_with(default)
209    }
210
211    fn or_insert_with_key<F>(self, default: F) -> &'a mut Self::Value
212    where
213        F: FnOnce(&Self::Key) -> Self::Value,
214    {
215        self.or_insert_with_key(default)
216    }
217
218    fn key(&self) -> &Self::Key {
219        self.key()
220    }
221
222    fn and_modify<F>(self, f: F) -> Self
223    where
224        F: FnOnce(&mut Self::Value),
225    {
226        self.and_modify(f)
227    }
228}
229
230#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
231pub struct HashMapFactory<K, V> {
232    _marker: PhantomData<fn() -> (K, V)>,
233}
234
235impl<K, V> Default for HashMapFactory<K, V> {
236    fn default() -> Self {
237        Self {
238            _marker: Default::default(),
239        }
240    }
241}
242
243#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
244pub struct HashMapFactoryWithCapacity<K, V> {
245    capacity: usize,
246    _marker: PhantomData<fn() -> (K, V)>,
247}
248
249impl<K, V> Default for HashMapFactoryWithCapacity<K, V> {
250    fn default() -> Self {
251        Self {
252            capacity: Default::default(),
253            _marker: Default::default(),
254        }
255    }
256}
257
258impl<K, V> ContainerFactory for HashMapFactory<K, V>
259where
260    K: Eq + Hash,
261{
262    type Container = HashMap<K, V>;
263
264    fn create_container(&self) -> Self::Container {
265        HashMap::new()
266    }
267}
268
269impl<K, V> ContainerFactory for HashMapFactoryWithCapacity<K, V>
270where
271    K: Eq + Hash,
272{
273    type Container = HashMap<K, V>;
274
275    fn create_container(&self) -> Self::Container {
276        HashMap::with_capacity(self.capacity)
277    }
278}
279
280impl<K, V> Container for HashMap<K, V>
281where
282    K: Eq + Hash,
283{
284    type Key = K;
285    type Value = V;
286    type Entry<'a>
287        = hash_map::Entry<'a, K, V>
288    where
289        K: 'a,
290        V: 'a;
291    type Iter<'a>
292        = hash_map::Iter<'a, K, V>
293    where
294        K: 'a,
295        V: 'a;
296    type Drain<'a>
297        = hash_map::Drain<'a, K, V>
298    where
299        K: 'a,
300        V: 'a;
301
302    fn get(&self, key: &Self::Key) -> Option<&Self::Value> {
303        self.get(key)
304    }
305
306    fn get_mut(&mut self, key: &Self::Key) -> Option<&mut Self::Value> {
307        self.get_mut(key)
308    }
309
310    fn insert(&mut self, key: Self::Key, value: Self::Value) -> Option<Self::Value> {
311        self.insert(key, value)
312    }
313
314    fn remove(&mut self, key: &Self::Key) -> Option<Self::Value> {
315        self.remove(key)
316    }
317
318    fn entry(&mut self, key: Self::Key) -> Self::Entry<'_> {
319        self.entry(key)
320    }
321
322    fn iter(&self) -> Self::Iter<'_> {
323        self.iter()
324    }
325
326    fn drain(&mut self) -> Self::Drain<'_> {
327        self.drain()
328    }
329}
330
331impl<'a, K, V> ContainerEntry<'a> for hash_map::Entry<'a, K, V>
332where
333    K: 'a + Eq + Hash,
334    V: 'a,
335{
336    type Key = K;
337    type Value = V;
338
339    fn or_default(self) -> &'a mut Self::Value
340    where
341        Self::Value: Default,
342    {
343        self.or_default()
344    }
345
346    fn or_insert(self, default: Self::Value) -> &'a mut Self::Value {
347        self.or_insert(default)
348    }
349
350    fn or_insert_with<F>(self, default: F) -> &'a mut Self::Value
351    where
352        F: FnOnce() -> Self::Value,
353    {
354        self.or_insert_with(default)
355    }
356
357    fn or_insert_with_key<F>(self, default: F) -> &'a mut Self::Value
358    where
359        F: FnOnce(&Self::Key) -> Self::Value,
360    {
361        self.or_insert_with_key(default)
362    }
363
364    fn key(&self) -> &Self::Key {
365        self.key()
366    }
367
368    fn and_modify<F>(self, f: F) -> Self
369    where
370        F: FnOnce(&mut Self::Value),
371    {
372        self.and_modify(f)
373    }
374}