competitive/data_structure/
vec_map.rs

1use super::{Container, ContainerEntry, ContainerFactory};
2use std::{
3    fmt::{self, Debug},
4    iter::{FilterMap, Map, repeat_with},
5    marker::PhantomData,
6    mem::replace,
7    slice,
8};
9
10#[derive(Debug, Clone)]
11pub struct VecMapFactory<K, V, F> {
12    key_to_index: F,
13    _marker: PhantomData<fn() -> (K, V)>,
14}
15
16#[derive(Debug, Clone)]
17pub struct FixedVecMapFactory<K, V, F> {
18    key_to_index: F,
19    len: usize,
20    _marker: PhantomData<fn() -> (K, V)>,
21}
22
23#[derive(Debug, Clone)]
24pub struct VecMapFactoryWithCapacity<K, V, F> {
25    key_to_index: F,
26    capacity: usize,
27    _marker: PhantomData<fn() -> (K, V)>,
28}
29
30#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
31pub struct VecMap<const FIXED: bool, K, V, F> {
32    pub data: Vec<Option<(K, V)>>,
33    key_to_index: F,
34}
35
36impl<const FIXED: bool, K, V, F> Debug for VecMap<FIXED, K, V, F>
37where
38    K: Debug,
39    V: Debug,
40{
41    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
42        f.debug_map()
43            .entries(self.data.iter().flatten().map(|(k, v): &(K, V)| (k, v)))
44            .finish()
45    }
46}
47
48#[derive(Debug)]
49pub struct Entry<'a, K, V> {
50    key: K,
51    entry: &'a mut Option<(K, V)>,
52}
53
54impl<K, V, F> VecMapFactory<K, V, F> {
55    pub fn new(key_to_index: F) -> Self {
56        Self {
57            key_to_index,
58            _marker: PhantomData,
59        }
60    }
61}
62
63impl<K, V, F> ContainerFactory for VecMapFactory<K, V, F>
64where
65    F: Fn(&K) -> usize + Clone,
66{
67    type Container = VecMap<false, K, V, F>;
68
69    fn create_container(&self) -> Self::Container {
70        VecMap {
71            data: Vec::new(),
72            key_to_index: self.key_to_index.clone(),
73        }
74    }
75}
76
77impl<K, V, F> FixedVecMapFactory<K, V, F> {
78    pub fn new(key_to_index: F, len: usize) -> Self {
79        Self {
80            key_to_index,
81            len,
82            _marker: PhantomData,
83        }
84    }
85}
86
87impl<K, V, F> ContainerFactory for FixedVecMapFactory<K, V, F>
88where
89    F: Fn(&K) -> usize + Clone,
90{
91    type Container = VecMap<true, K, V, F>;
92
93    fn create_container(&self) -> Self::Container {
94        VecMap {
95            data: repeat_with(|| None).take(self.len).collect(),
96            key_to_index: self.key_to_index.clone(),
97        }
98    }
99}
100
101impl<K, V, F> VecMapFactoryWithCapacity<K, V, F> {
102    pub fn new(key_to_index: F, capacity: usize) -> Self {
103        Self {
104            key_to_index,
105            capacity,
106            _marker: PhantomData,
107        }
108    }
109}
110
111impl<K, V, F> ContainerFactory for VecMapFactoryWithCapacity<K, V, F>
112where
113    F: Fn(&K) -> usize + Clone,
114{
115    type Container = VecMap<false, K, V, F>;
116
117    fn create_container(&self) -> Self::Container {
118        VecMap {
119            data: Vec::with_capacity(self.capacity),
120            key_to_index: self.key_to_index.clone(),
121        }
122    }
123}
124
125impl<const FIXED: bool, K, V, F> Container for VecMap<FIXED, K, V, F>
126where
127    F: Fn(&K) -> usize,
128{
129    type Key = K;
130    type Value = V;
131    type Entry<'a>
132        = Entry<'a, K, V>
133    where
134        Self: 'a,
135        Self::Key: 'a,
136        Self::Value: 'a;
137    type Iter<'a>
138        = Map<
139        FilterMap<slice::Iter<'a, Option<(K, V)>>, fn(&Option<(K, V)>) -> Option<&(K, V)>>,
140        fn(&(K, V)) -> (&K, &V),
141    >
142    where
143        Self: 'a,
144        Self::Key: 'a,
145        Self::Value: 'a;
146    type Drain<'a>
147        = FilterMap<slice::IterMut<'a, Option<(K, V)>>, fn(&mut Option<(K, V)>) -> Option<(K, V)>>
148    where
149        Self: 'a,
150        Self::Key: 'a,
151        Self::Value: 'a;
152
153    fn get(&self, key: &Self::Key) -> Option<&Self::Value> {
154        let index = (self.key_to_index)(key);
155        self.data.get(index).and_then(|x| x.as_ref()).map(|x| &x.1)
156    }
157
158    fn get_mut(&mut self, key: &Self::Key) -> Option<&mut Self::Value> {
159        let index = (self.key_to_index)(key);
160        self.data
161            .get_mut(index)
162            .and_then(|x| x.as_mut())
163            .map(|x| &mut x.1)
164    }
165
166    fn insert(&mut self, key: Self::Key, value: Self::Value) -> Option<Self::Value> {
167        let index = (self.key_to_index)(&key);
168        let entry = if FIXED {
169            self.data.get_mut(index).unwrap()
170        } else {
171            if index >= self.data.len() {
172                self.data.resize_with(index + 1, Default::default);
173            }
174            unsafe { self.data.get_unchecked_mut(index) }
175        };
176        match entry {
177            Some((_, v)) => Some(replace(v, value)),
178            entry => {
179                *entry = Some((key, value));
180                None
181            }
182        }
183    }
184
185    fn remove(&mut self, key: &Self::Key) -> Option<Self::Value> {
186        let index = (self.key_to_index)(key);
187        self.data.get_mut(index).and_then(|x| x.take()).map(|x| x.1)
188    }
189
190    fn entry(&mut self, key: Self::Key) -> Self::Entry<'_> {
191        let index = (self.key_to_index)(&key);
192        let entry = if FIXED {
193            self.data.get_mut(index).unwrap()
194        } else {
195            if index >= self.data.len() {
196                self.data.resize_with(index + 1, Default::default);
197            }
198            unsafe { self.data.get_unchecked_mut(index) }
199        };
200        Entry { key, entry }
201    }
202
203    fn iter(&self) -> Self::Iter<'_> {
204        self.data
205            .iter()
206            .filter_map::<_, fn(&Option<(K, V)>) -> Option<&(K, V)>>(Option::as_ref)
207            .map::<_, fn(&(K, V)) -> (&K, &V)>(|(k, v)| (k, v))
208    }
209
210    fn drain(&mut self) -> Self::Drain<'_> {
211        self.data
212            .iter_mut()
213            .filter_map::<_, fn(&mut Option<(K, V)>) -> Option<(K, V)>>(Option::take)
214    }
215}
216
217impl<'a, K, V> ContainerEntry<'a> for Entry<'a, K, V> {
218    type Key = K;
219    type Value = V;
220
221    fn or_default(self) -> &'a mut Self::Value
222    where
223        Self::Value: Default,
224    {
225        match self.entry {
226            Some((_, value)) => value,
227            entry => {
228                *entry = Some((self.key, Default::default()));
229                unsafe { &mut entry.as_mut().unwrap_unchecked().1 }
230            }
231        }
232    }
233
234    fn or_insert(self, default: Self::Value) -> &'a mut Self::Value {
235        match self.entry {
236            Some((_, value)) => value,
237            entry => {
238                *entry = Some((self.key, default));
239                unsafe { &mut entry.as_mut().unwrap_unchecked().1 }
240            }
241        }
242    }
243
244    fn or_insert_with<F>(self, default: F) -> &'a mut Self::Value
245    where
246        F: FnOnce() -> Self::Value,
247    {
248        match self.entry {
249            Some((_, value)) => value,
250            entry => {
251                *entry = Some((self.key, default()));
252                unsafe { &mut entry.as_mut().unwrap_unchecked().1 }
253            }
254        }
255    }
256
257    fn or_insert_with_key<F>(self, default: F) -> &'a mut Self::Value
258    where
259        F: FnOnce(&Self::Key) -> Self::Value,
260    {
261        match self.entry {
262            Some((_, value)) => value,
263            entry => {
264                let default = default(&self.key);
265                *entry = Some((self.key, default));
266                unsafe { &mut entry.as_mut().unwrap_unchecked().1 }
267            }
268        }
269    }
270
271    fn key(&self) -> &Self::Key {
272        &self.key
273    }
274
275    fn and_modify<F>(self, f: F) -> Self
276    where
277        F: FnOnce(&mut Self::Value),
278    {
279        if let Some((_, value)) = self.entry {
280            f(value);
281        }
282        self
283    }
284}
285
286#[cfg(test)]
287mod tests {
288    use super::*;
289
290    #[test]
291    fn test_vec_map_get() {
292        let mut map = VecMapFactory::new(|&x: &usize| x).create_container();
293        map.insert(0, 0);
294        assert_eq!(map.get(&0), Some(&0));
295        assert_eq!(map.get(&1), None);
296    }
297
298    #[test]
299    fn test_vec_map_get_mut() {
300        let mut map = VecMapFactory::new(|&x: &usize| x).create_container();
301        map.insert(0, 0);
302        *map.get_mut(&0).unwrap() += 1;
303        assert_eq!(map.get(&0), Some(&1));
304    }
305
306    #[test]
307    fn test_vec_map_insert() {
308        let mut map = VecMapFactory::new(|&x: &usize| x).create_container();
309        assert_eq!(map.insert(0, 0), None);
310        assert_eq!(map.insert(0, 1), Some(0));
311        assert_eq!(map.get(&0), Some(&1));
312    }
313
314    #[test]
315    fn test_vec_map_remove() {
316        let mut map = VecMapFactory::new(|&x: &usize| x).create_container();
317        map.insert(0, 0);
318        assert_eq!(map.remove(&0), Some(0));
319        assert_eq!(map.remove(&0), None);
320    }
321
322    #[test]
323    fn test_vec_map_entry() {
324        let mut map = VecMapFactory::new(|&x: &usize| x).create_container();
325        map.entry(0).or_insert(0);
326        assert_eq!(*map.entry(0).or_insert(1), 0);
327        assert_eq!(*map.entry(1).and_modify(|e| *e += 1).or_default(), 0);
328        assert_eq!(*map.entry(1).and_modify(|e| *e += 1).or_default(), 1);
329    }
330}
competitive/data_structure/vec_map.rs

competitive/data_structure/
vec_map.rs
