competitive/data_structure/
container.rs1use 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}